import numpy as np
import math


def getKref(k, rank, pn, ref):

    Nz = k.shape[0]
    nz = Nz // pn
    if ref == 1:
        return getK(k, rank, pn)

    if (rank > 0) and (rank < pn - 1):

        k = k[rank * nz - 1 : (rank + 1) * nz + 1, :, :]
        k = refinaPy(k, ref)
        if ref != 1:
            k = k[(ref - 1) : -(ref - 1), :, :]
        nz, ny, nx = k.shape[0], k.shape[1], k.shape[2]
        ki = np.zeros((nz, ny + 2, nx + 2))
        ki[:, 1:-1, 1:-1] = k
        # print(ki.shape)
        nnz = nz - 2
    if rank == 0:
        k = k[: (rank + 1) * nz + 1, :, :]
        k = refinaPy(k, ref)
        if ref != 1:
            k = k[: -(ref - 1), :, :]
        nz, ny, nx = k.shape[0], k.shape[1], k.shape[2]
        ki = np.zeros((nz + 1, ny + 2, nx + 2))
        ki[1:, 1:-1, 1:-1] = k
        ki[0, :, :] = ki[1, :, :]
        nnz = nz
    if rank == (pn - 1):
        k = k[rank * nz - 1 :, :, :]
        k = refinaPy(k, ref)
        if ref != 1:
            k = k[(ref - 1) :, :, :]
        nz, ny, nx = k.shape[0], k.shape[1], k.shape[2]
        ki = np.zeros((nz + 1, ny + 2, nx + 2))
        ki[:-1, 1:-1, 1:-1] = k
        ki[-1, :, :] = ki[-2, :, :]
        nnz = (Nz // pn) * ref
    return ki, Nz * ref, nnz


def getK(k, rank, pn):

    # k=np.load(kfile)
    # nn=int(np.cbrt(k.shape[0]))
    # k=k.reshape((nn,nn,nn))

    Nz, Ny, Nx = k.shape[0], k.shape[1], k.shape[2]
    nz = Nz // pn
    if rank == pn - 1:
        nnz = Nz - (pn - 1) * nz
        ki = np.zeros((nnz + 2, Ny + 2, Nx + 2))
    else:
        nnz = nz
        ki = np.zeros((nz + 2, Ny + 2, Nx + 2))
    if (rank > 0) and (rank < pn - 1):
        ki[:, 1:-1, 1:-1] = k[rank * nz - 1 : (rank + 1) * nz + 1, :, :]
    if rank == 0:
        ki[1:, 1:-1, 1:-1] = k[: (rank + 1) * nz + 1, :, :]
        ki[0, :, :] = ki[1, :, :]
    if rank == (pn - 1):
        ki[:-1, 1:-1, 1:-1] = k[rank * nz - 1 :, :, :]
        ki[-1, :, :] = ki[-2, :, :]
    return ki, Nz, nz


"""
def getK(k,rank,pn):

	#k=np.load(kfile)
	#nn=int(np.cbrt(k.shape[0]))
	#k=k.reshape((nn,nn,nn))
	
	Nz, Ny,Nx=k.shape[0],k.shape[1],k.shape[2]
	nz=Nz//pn  
	if rank==pn-1:
		nnz= Nz-(pn-1)*nz
		ki=np.zeros((nnz+2,Ny+2,Nx+2))
	else:
		nnz=nz
		ki=np.zeros((nz+2,Ny+2,Nx+2))
	if (rank>0) and (rank<pn-1):
		ki[:,1:-1,1:-1]=k[rank*nz-1:(rank+1)*nz+1,:,:]
	if rank==0:
		ki[1:,1:-1,1:-1]=k[:(rank+1)*nz+1,:,:]
		ki[0,:,:]=ki[1,:,:]
	if rank==(pn-1):
		ki[:-1,1:-1,1:-1]=k[rank*nz-1:,:,:]
		ki[-1,:,:]=ki[-2,:,:]
	return ki, Nz, nz
"""


def refinaPy(k, ref):

    if ref == 1:
        return k

    nx, ny, nz = k.shape[2], k.shape[1], k.shape[0]

    krz = np.zeros((ref * nz, ny, nx))
    for i in range(ref):
        krz[i::ref, :, :] = k
    k = 0
    krzy = np.zeros((ref * nz, ny * ref, nx))
    for i in range(ref):
        krzy[:, i::ref, :] = krz

    if nx == 1:
        return krzy
    krz = 0

    krzyx = np.zeros((ref * nz, ny * ref, nx * ref))
    for i in range(ref):
        krzyx[:, :, i::ref] = krzy

    return krzyx  # krzyx[(ref-1):-(ref-1),:,:]


def centL(K, R, kkm, r):

    nx, ny, nz = kkm.shape[2] - 2, kkm.shape[1] - 2, kkm.shape[0] - 2
    for k in range(nz):
        for j in range(ny):
            for i in range(nx):
                t = np.array(
                    [
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 1, i + 2]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i + 2]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 1, i]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 2, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 2, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 2, j + 1, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 2, j + 1, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k, j + 1, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k, j + 1, i + 1]),
                    ]
                )

                K.setValues(r, r, t[0] + t[1] + t[2] + t[3] + t[4] + t[5])
                K.setValues(r, r + 1, -t[0])
                K.setValues(r, r - 1, -t[1])
                K.setValues(r, r + nx, -t[2])
                K.setValues(r, r - nx, -t[3])
                K.setValues(r, r + nx * ny, -t[4])
                K.setValues(r, r - nx * ny, -t[5])
                R.setValues(r, 0)
                r += 1
    return K, R


def firstL(K, R, kkm, pbc):
    # Right side of Rmat
    r = 0

    nx, ny, nz = kkm.shape[2] - 2, kkm.shape[1] - 2, kkm.shape[0] - 2

    k = 0

    for j in range(ny):
        for i in range(nx):
            t = np.array(
                [
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j + 1, i + 2]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i + 2]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j + 1, i]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j + 2, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 2, i + 1]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j, i + 1]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 2, j + 1, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 2, j + 1, i + 1]),
                    4
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k, j + 1, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k, j + 1, i + 1]),
                ]
            )  # atento aca BC 2Tz
            K.setValues(r, r, t[0] + t[1] + t[2] + t[3] + t[4] + t[5])
            K.setValues(r, r + 1, -t[0])
            K.setValues(r, r + nx, -t[2])
            K.setValues(r, r + nx * ny, -t[4])
            R.setValues(r, t[5] * pbc)
            r += 1

        # Left side of Rmat
    for j in range(ny):
        for i in range(1, nx):
            r = j * nx + i
            K.setValues(
                r,
                r - 1,
                -2
                * kkm[k + 1, j + 1, i + 1]
                * kkm[k + 1, j + 1, i]
                / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i]),
            )

    for j in range(1, ny):
        for i in range(nx):
            r = j * nx + i
            K.setValues(
                r,
                r - nx,
                -2
                * kkm[k + 1, j + 1, i + 1]
                * kkm[k + 1, j, i + 1]
                / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j, i + 1]),
            )

    for k in range(1, nz):
        for j in range(ny):
            for i in range(nx):
                r = k * ny * nx + j * nx + i
                t = np.array(
                    [
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 1, i + 2]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i + 2]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 1, i]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 2, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 2, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 2, j + 1, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 2, j + 1, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k, j + 1, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k, j + 1, i + 1]),
                    ]
                )
                K.setValues(r, r, t[0] + t[1] + t[2] + t[3] + t[4] + t[5])
                K.setValues(r, r + 1, -t[0])
                K.setValues(r, r - 1, -t[1])
                K.setValues(r, r + nx, -t[2])
                K.setValues(r, r - nx, -t[3])
                K.setValues(r, r + nx * ny, -t[4])
                K.setValues(r, r - nx * ny, -t[5])
                R.setValues(r, 0)
    return K, R


def lastL(K, R, kkm, r):
    # Right side of Rmat

    nx, ny, nz = kkm.shape[2] - 2, kkm.shape[1] - 2, kkm.shape[0] - 2

    for k in range(nz - 1):
        for j in range(ny):
            for i in range(nx):
                t = np.array(
                    [
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 1, i + 2]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i + 2]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 1, i]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j + 2, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 2, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 1, j, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k + 2, j + 1, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k + 2, j + 1, i + 1]),
                        2
                        * kkm[k + 1, j + 1, i + 1]
                        * kkm[k, j + 1, i + 1]
                        / (kkm[k + 1, j + 1, i + 1] + kkm[k, j + 1, i + 1]),
                    ]
                )
                K.setValues(r, r, t[0] + t[1] + t[2] + t[3] + t[4] + t[5])
                K.setValues(r, r + 1, -t[0])
                K.setValues(r, r - 1, -t[1])
                K.setValues(r, r + nx, -t[2])
                K.setValues(r, r - nx, -t[3])
                K.setValues(r, r + nx * ny, -t[4])
                K.setValues(r, r - nx * ny, -t[5])
                R.setValues(r, 0)
                r = r + 1

    auxr = r
    k = -3

    for j in range(ny):
        for i in range(nx):

            t = np.array(
                [
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j + 1, i + 2]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i + 2]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j + 1, i]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j + 2, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 2, i + 1]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 1, j, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j, i + 1]),
                    4
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k + 2, j + 1, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k + 2, j + 1, i + 1]),
                    2
                    * kkm[k + 1, j + 1, i + 1]
                    * kkm[k, j + 1, i + 1]
                    / (kkm[k + 1, j + 1, i + 1] + kkm[k, j + 1, i + 1]),
                ]
            )  # guarda aca BC en t[4] va por 2 por dx/2

            K.setValues(r, r, t[0] + t[1] + t[2] + t[3] + t[4] + t[5])
            K.setValues(r, r - 1, -t[1])
            K.setValues(r, r - nx, -t[3])
            K.setValues(r, r - nx * ny, -t[5])
            R.setValues(r, 0)
            r += 1

        # Right side of Mat
    for j in range(ny):
        for i in range(nx - 1):
            r = j * nx + i + auxr
            K.setValues(
                r,
                r + 1,
                -2
                * kkm[k + 1, j + 1, i + 1]
                * kkm[k + 1, j + 1, i + 2]
                / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 1, i + 2]),
            )

    for j in range(ny - 1):
        for i in range(nx):
            r = j * nx + i + auxr
            K.setValues(
                r,
                r + nx,
                -2
                * kkm[k + 1, j + 1, i + 1]
                * kkm[k + 1, j + 2, i + 1]
                / (kkm[k + 1, j + 1, i + 1] + kkm[k + 1, j + 2, i + 1]),
            )

    return K, R