import numpy as np
import os
import time
from tools.postprocessK.flow import ComputeVol, comp_Kdiss_Kaverage
import subprocess

#k[x,y,z]
import json

def comp_postKeff(parser,rundir,nr,PetscP):


	k=np.load(rundir+'k.npy')
	P=np.load(rundir+'P.npy')
	ref=P.shape[0]//k.shape[0]

	t0=time.time()
	k, diss, vx, Px, Py, Pz = ComputeVol(k,P) #refina k
	tDissVel=time.time()-t0

	P=0


	S_min_post = int(parser.get('K-Postprocess','MinBlockSize'))
	nimax =2** int(parser.get('K-Postprocess','Max_sample_size'))
	compKperm =parser.get('K-Postprocess','kperm')
	if compKperm=='yes':
		compKperm=True		

	S_min_post=S_min_post*ref

	if S_min_post==0:
		sx=k.shape[0]
	else:
		sx = get_min_nbl(k,nimax,nr,S_min_post)  


	kdiss,kave=getKpost(k, diss, vx, Px, Py, Pz,sx,rundir,ref)


	ttotal=time.time()-t0

	summary = np.array([kdiss,kave,ttotal,tDissVel/ttotal]).T
	np.savetxt(rundir + 'PosKeffSummary.txt',summary,header='K_diss, K_average,ttotal,tDiss/ttotal')

	return 




def getKpost(kf, diss, vx, Px, Py, Pz,sx,rundir,ref,compkperm):



	ex=int(np.log2(kf.shape[0]))
	esx=int(np.log2(sx))

	scales=2**np.arange(esx,ex)

	datadir=rundir+'KpostProcess/'
	try:
		os.makedirs(datadir)
	except:
		nada=0

	for l in scales:
		nblx, nbly, nblz = kf.shape[0]//l, kf.shape[1]//l, kf.shape[2]//l
		sx,sy,sz=l,l,l		
		if kf.shape[2]==1:
			nblz=1
			sz=1
		

		Kdiss,Kave=np.zeros((nblx,nbly,nblz)),np.zeros((nblx,nbly,nblz))
		if compkperm==True:
			Kperm = np.zeros((nblx,nbly,nblz))

		for i in range(nblx):
			for j in range(nbly):
				for k in range(nblz):
					Kdiss[i,j,k],Kave[i,j,k]=comp_Kdiss_Kaverage(kf[i*sx:(i+1)*sx,j*sy:(j+1)*sy,k*sz:(k+1)*sz], diss[i*sx:(i+1)*sx,j*sy:(j+1)*sy,k*sz:(k+1)*sz], vx[i*sx:(i+1)*sx,j*sy:(j+1)*sy,k*sz:(k+1)*sz], Px[i*sx:(i+1)*sx+1,j*sy:(j+1)*sy+1,k*sz:(k+1)*sz+1], Py[i*sx:(i+1)*sx+1,j*sy:(j+1)*sy+1,k*sz:(k+1)*sz+1], Pz[i*sx:(i+1)*sx+1,j*sy:(j+1)*sy+1,k*sz:(k+1)*sz+1])
					if compkperm==True:
						Kperm[i,j,k]=PetscP(datadir,ref,k)(kf[i*sx:(i+1)*sx,j*sy:(j+1)*sy,k*sz:(k+1)*sz])
					


		np.save(datadir+'Kd'+str(l//ref)+'.npy',Kdiss)
		np.save(datadir+'Kv'+str(l//ref)+'.npy',Kave)
		if compkperm==True:
			np.save(datadir+'Kperm'+str(l//ref)+'.npy',Kperm)

	Kdiss,Kave = comp_Kdiss_Kaverage(kf, diss, vx, Px, Py, Pz)
	
	np.save(datadir+'Kd'+str(kf.shape[0]//ref)+'.npy',np.array([Kdiss]))
	np.save(datadir+'Kv'+str(kf.shape[0]//ref)+'.npy',np.array([Kave]))


	return Kdiss, Kave


def get_min_nbl(kc,nimax,nr,smin):

	if kc.shape[2]==1:
		dim=2.0
	else:
		dim=3.0
	if nr>0:
		y=(1/dim)*np.log2(nr*kc.size/(nimax*(smin**dim)))
	else:	
		y=0
	y=int(y)
	s=int((2**y) * smin)
	if s<smin:
		s=smin

	return s