import numpy as np
import sys
from refine import refine as ref

def get_p(pn, pdir, pprefix):


	p=np.load(pdir+pprefix+"0"+'.npy')
	for i in range(1,pn):
		p=np.concatenate((p,np.load(pdir+pprefix+str(i)+'.npy')),axis=0)

	return p



def get_k(pn, kdir, kprefix):

	k=(np.load(kdir+kprefix+'0'+'.npy'))[1:-1,:,:]
	for i in range(1,pn):
		k=np.concatenate((k,(np.load(kdir+kprefix+str(i)+'.npy'))[1:-1,:,:]),axis=0)
	return ref(k,2,2,2)






def kef(P,K,i,j,k,pbc):
	#tx=2*K[:,:,i]*K[:,:,i+1]/(K[:,:,i]+K[:,:,i+1])
	#ty=2*K[:,j,:]*K[:,j+1,:]/(K[:,j,:]+K[:,j+1,:])
	tz=2*K[k,:,:]*K[k+1,:,:]/(K[k,:,:]+K[k+1,:,:])

	#qx=tx*(P[:,:,i+1]-P[:,:,i])
	#qy=ty*(P[:,j+1,:]-P[:,j,:])
	qz=-tz*(P[k+1,:,:]-P[k,:,:])

	kz=qz.sum()*(K.shape[0]+1)/(pbc*K.shape[1]*K.shape[2])
	return kz


def test(pn, kdir, pdir, kprefix, pprefix):

	K=get_k(pn, kdir, kprefix)
	print(K.shape)
	P=get_p(pn, pdir, pprefix)
	print(P)
	print(P.shape)
	print(kef(P,K,1,1,1,1000))
	return


pn=4
kdir="./test/"
pdir="./test/"
kprefix="k"
pprefix="P"


test(pn, kdir, pdir, kprefix, pprefix)