import numpy as np

from scipy.sparse import diags
from scipy.stats import mstats
from scipy.sparse.linalg import (
    bicg,
    bicgstab,
    cg,
    dsolve,
)  # ,LinearOperator, spilu, bicgstab

# from scikits.umfpack import spsolve, splu

import time


def getDiss(k, vx, vy, vz):
    diss = (
        vx[1:, :, :] ** 2
        + vx[:-1, :, :] ** 2
        + vy[:, 1:, :] ** 2
        + vy[:, :-1, :] ** 2
        + vz[:, :, 1:] ** 2
        + vz[:, :, :-1] ** 2
    ) / (2 * k)
    return diss


def ComputeVol(k, P, saveV):

    k = refina(k, P.shape[0] // k.shape[0])
    Px, Py, Pz = getPfaces(k, P)
    vx, vy, vz = getVfaces(k, P, Px, Py, Pz)
    diss = getDiss(k, vx, vy, vz)
    if saveV == False:
        vy, vz = 0, 0
    else:
        vy, vz = 0.5 * (vy[:, 1:, :] + vy[:, :-1, :]), 0.5 * (
            vz[:, :, 1:] + vz[:, :, :-1]
        )
    vx = 0.5 * (vx[1:, :, :] + vx[:-1, :, :])

    return k, diss, vx, vy, vz, Px, Py, Pz


def comp_Kdiss_Kaverage(k, diss, vx, Px, Py, Pz):

    mgx, mgy, mgz = (
        np.mean(Px[-1, :, :] - Px[0, :, :]) / k.shape[0],
        np.mean(Py[:, -1, :] - Py[:, 0, :]) / k.shape[1],
        np.mean(Pz[:, :, -1] - Pz[:, :, 0]) / k.shape[2],
    )
    kave = np.mean(vx) / mgx
    kdiss = np.mean(diss) / (mgx ** 2 + mgy ** 2 + mgz ** 2)
    return kdiss, kave


def getKeff(pm, k, pbc, Nz):

    nx = k.shape[2]  # Pasar k sin bordes de k=0
    ny = k.shape[1]

    tz = 2 * k[1, :, :] * k[0, :, :] / (k[0, :, :] + k[1, :, :])
    q = ((pm[0, :, :] - pm[1, :, :]) * tz).sum()
    area = ny * nx
    l = Nz
    keff = q * l / (pbc * area)
    return keff, q


def getPfaces(k, P):
    nx, ny, nz = k.shape[0], k.shape[1], k.shape[2]
    Px, Py, Pz = (
        np.zeros((nx + 1, ny, nz)),
        np.zeros((nx, ny + 1, nz)),
        np.zeros((nx, ny, nz + 1)),
    )

    Px[1:-1, :, :] = (k[:-1, :, :] * P[:-1, :, :] + k[1:, :, :] * P[1:, :, :]) / (
        k[:-1, :, :] + k[1:, :, :]
    )
    Px[0, :, :] = nx

    Py[:, 1:-1, :] = (k[:, :-1, :] * P[:, :-1, :] + k[:, 1:, :] * P[:, 1:, :]) / (
        k[:, :-1, :] + k[:, 1:, :]
    )
    Py[:, 0, :], Py[:, -1, :] = P[:, 0, :], P[:, -1, :]

    Pz[:, :, 1:-1] = (k[:, :, :-1] * P[:, :, :-1] + k[:, :, 1:] * P[:, :, 1:]) / (
        k[:, :, :-1] + k[:, :, 1:]
    )
    Pz[:, :, 0], Pz[:, :, -1] = P[:, :, 0], P[:, :, -1]

    return Px, Py, Pz


def getVfaces(k, P, Px, Py, Pz):
    nx, ny, nz = k.shape[0], k.shape[1], k.shape[2]
    vx, vy, vz = (
        np.zeros((nx + 1, ny, nz)),
        np.zeros((nx, ny + 1, nz)),
        np.zeros((nx, ny, nz + 1)),
    )
    vx[1:, :, :] = 2 * k * (Px[1:, :, :] - P)  # v= k*(deltaP)/(deltaX/2)
    vx[0, :, :] = 2 * k[0, :, :] * (P[0, :, :] - Px[0, :, :])

    vy[:, 1:, :] = 2 * k * (Py[:, 1:, :] - P)
    vy[:, 0, :] = 2 * k[:, 0, :] * (P[:, 0, :] - Py[:, 0, :])
    vz[:, :, 1:] = 2 * k * (Pz[:, :, 1:] - P)
    vz[:, :, 0] = 2 * k[:, :, 0] * (P[:, :, 0] - Pz[:, :, 0])
    return vx, vy, vz


def refina(k, ref):

    if ref == 1:
        return k

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

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

    for i in range(ref):
        krxy[:, i::ref, :] = krx
    krx = 0
    if nz == 1:
        return krxy

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

    return krxyz


def computeT(k):

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

    return tx, ty, tz


def Rmat(k):

    pbc = k.shape[0]
    tx, ty, tz = computeT(k)

    tx[0, :, :], tx[-1, :, :] = 2 * tx[1, :, :], 2 * tx[-2, :, :]

    rh = np.zeros((k.shape[0], k.shape[1], k.shape[2]))

    rh[0, :, :] = pbc * tx[0, :, :]
    rh = rh.reshape(-1)
    d = (
        tz[:, :, :-1]
        + tz[:, :, 1:]
        + ty[:, :-1, :]
        + ty[:, 1:, :]
        + tx[:-1, :, :]
        + tx[1:, :, :]
    ).reshape(-1)
    a = (-tz[:, :, :-1].reshape(-1))[1:]
    # a=(tx.reshape(-1))[:-1]
    b = (-ty[:, 1:, :].reshape(-1))[: -k.shape[2]]
    c = -tx[1:-1, :, :].reshape(-1)

    return a, b, c, d, rh


def PysolveP(k, solver):
    a, b, c, d, rh = Rmat(k)
    nx, ny, nz = k.shape[0], k.shape[1], k.shape[2]
    offset = [-nz * ny, -nz, -1, 0, 1, nz, nz * ny]
    km = diags(np.array([c, b, a, d, a, b, c]), offset, format="csc")
    a, b, c, d = 0, 0, 0, 0
    lu = splu(km)
    print(lu)
    p = solver(km, rh)
    p = p.reshape(nx, ny, nz)
    keff, q = getKeff(p, k, nz, nz)
    return keff


"""
solvers=[bicg, bicgstab, cg, dsolve, spsolve]
snames=['bicg', 'bicgstab',' cg',' dsolve',' spsolve']

for job in range(jobs):
	kff=np.load('./otrotest/'+str(job)+'/k.npy')
	print('*************   JOB :    '+str(job)+'   ******************')
	print('   ')
	for i in range(len(solvers)):

		t0=time.time()


		keff=PysolveP(kff, solvers[i])
		print('Solver:  '+snames[i]+'   time:   '+str(time.time()-t0))
"""