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