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