import numpy as np
from scipy.sparse import diags
from scipy.stats import mstats
from scipy.sparse.linalg import (
    spsolve,
    bicg,
    bicgstab,
    cg,
)  # ,LinearOperator, spilu, bicgstab
from petsc4py import PETSc
import csv
import time

# [layer,columns,row]= [z,y,x]


NNN = 256
ref = 2


def computeT(k):

    nx = k.shape[2]
    ny = k.shape[1]
    nz = k.shape[0] - 2
    tx = np.zeros((nz, ny, nx + 1))
    ty = np.zeros((nz, ny + 1, nx))
    tz = np.zeros((nz + 1, ny, nx))
    tx[:, :, 1:-1] = (
        2 * k[1:-1, :, :-1] * k[1:-1, :, 1:] / (k[1:-1, :, :-1] + k[1:-1, :, 1:])
    )
    ty[:, 1:-1, :] = (
        2 * k[1:-1, :-1, :] * k[1:-1, 1:, :] / (k[1:-1, :-1, :] + k[1:-1, 1:, :])
    )
    tz[:, :, :] = 2 * k[:-1, :, :] * k[1:, :, :] / (k[:-1, :, :] + k[1:, :, :])

    return tx, ty, tz, nx, ny, nz


def rafina(k, ref):

    if ref == 1:
        return k

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

    return krzy


def get_kfield():

    # auxk=np.load('k.npy')
    # auxk=auxk.reshape(nz,ny,nx)
    # k=np.ones((nz,ny,nx))
    # k = np.random.lognormal(0,3,(nz,ny,nx))
    # k=np.load('./inp/k.npy')
    # N=512
    # k=np.loadtxt('./inp/out_rafine.dat')
    # n=int(np.sqrt(k.size))
    # k=k.reshape((n,n))
    # k=k[:N,:N]
    k = np.load("k.npy")
    # kfiledir='../Modflow/bin/r'+str(ref)+'/'
    # k=np.loadtxt(kfiledir+'out_fftma.txt')
    # k=np.loadtxt(kfiledir+'out_rafine.dat')
    # k=k.reshape(int(np.sqrt(k.size)),int(np.sqrt(k.size)),1)
    # k=k[:NNN*ref,:NNN*ref,:]
    # k=rafina(k,ref)
    nx, ny, nz = k.shape[2], k.shape[1], k.shape[0]
    # k=k.reshape((nz,ny,nx))
    auxk = np.zeros((nz + 2, ny, nx))
    auxk[1:-1, :, :] = k
    auxk[0, :, :] = k[0, :, :]
    auxk[-1, :, :] = k[-1, :, :]

    return auxk


def Rmat(k, pbc):

    tx, ty, tz, nx, ny, nz = computeT(k)

    rh = np.zeros((nz, ny, nx))
    rh[0, :, :] = pbc * tz[0, :, :]
    rh = rh.reshape(-1)
    d = (
        tx[:, :, :-1]
        + tx[:, :, 1:]
        + ty[:, :-1, :]
        + ty[:, 1:, :]
        + tz[:-1, :, :]
        + tz[1:, :, :]
    ).reshape(-1)
    a = (-tx[:, :, :-1].reshape(-1))[1:]
    # a=(tx.reshape(-1))[:-1]
    b = (-ty[:, 1:, :].reshape(-1))[:-nx]
    c = -tz[1:-1, :, :].reshape(-1)

    return a, b, c, d, rh


def imp(k):
    for i in range(k.shape[1]):
        for j in range(k.shape[0]):
            if k[j, i] != 0:
                print(i, j, k[j, i])
    return


def PysolveP(a, b, c, d, rh, nx, ny, nz, solver):
    offset = [-nx * ny, -nx, -1, 0, 1, nx, nx * ny]
    k = diags(np.array([c, b, a, d, a, b, c]), offset, format="csc")

    p = solver(k, rh)
    return p


def PysolveP2d(b, c, d, rh, nx, ny, nz, solver):
    offset = [-ny, -1, 0, 1, ny]
    k = diags(np.array([c, b, d, b, c]), offset, format="csc")
    # imp(k.toarray())
    p = solver(k, rh)
    return p


def Pmat(pm, nx, ny, nz, pbc):
    auxpm = np.zeros((nz + 2, ny, nx))
    auxpm[0, :, :] = pbc
    auxpm[1:-1, :, :] = pm.reshape(nz, ny, nx)
    return auxpm


def getK(pm, k, pbc):

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

    tz = 2 * k[2, :, :] * k[1, :, :] / (k[2, :, :] + k[1, :, :])
    q = ((pm[1, :, :] - pm[2, :, :]) * tz).sum()

    area = nx * ny
    l = nz + 1

    keff = q * l / (pbc * area)

    # print('Arit = ', np.mean(k),'  Geom = ',mstats.gmean(k,axis=None),'  Harm = ',mstats.hmean(k, axis=None))

    return keff


def main():
    pbc = 1000
    solver = spsolve

    k = get_kfield()

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

    # print(k.shape)
    a, b, c, d, rh = Rmat(k, pbc)
    if nx == 1:
        p = PysolveP2d(b, c, d, rh, nx, ny, nz, solver)
    else:
        p = PysolveP(a, b, c, d, rh, nx, ny, nz, solver)
    print(p.shape)

    p = Pmat(p, nx, ny, nz, pbc)
    p = p.reshape((nz + 2, ny, nx))
    keff = getK(p, k, pbc)
    print(keff)

    # k=k.reshape((nz+2,ny,nx))

    auxp = np.zeros((nz + 2, ny + 2))
    auxk = np.zeros((nz + 2, ny + 2))
    auxp[:, 0] = p[:, 0]
    auxp[:, -1] = p[:, -1]
    auxp[:, 1:-1] = p

    auxk[:, 0] = 0
    auxk[:, -1] = 0
    auxk[:, 1:-1] = k

    np.save("./p", auxp)
    np.save("./k", auxk)
    # np.savetxt('./1p/k.txt',auxk)
    np.savetxt("./keff.txt", np.array([keff]))
    # print(p)

    return


main()