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))
'''