import numpy as np
import os
import time
from tools.solver.Ndar import PetscP

def comp_kperm_sub(parser,rundir,nr):


	k=np.load(rundir+'k.npy')

	ref=int(parser.get('Solver',"ref"))


	t0=time.time()

	S_min_post = int(parser.get('K-Postprocess','MinBlockSize'))
	nimax =2** int(parser.get('K-Postprocess','Max_sample_size'))
		

	S_min_post=S_min_post*ref

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

	if sx==1:
		sx=2

	tkperm=getKpost(k, sx,rundir,ref)


	ttotal=time.time()-t0

	return 




def getKpost(kf, sx,rundir,ref):



	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

	tkperm=np.zeros((scales.shape[0]))

	for il in range(scales.shape[0]):
		l=scales[il]
		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
		if l==2:
			refDeg=2
		else:
			refDeg=ref


		tkperm[il]=time.time()
		Kperm = np.zeros((nblx,nbly,nblz))
		try:
			Kperm=np.load(datadir+'Kperm'+str(l//ref)+'.npy')
			
		except:
			for i in range(nblx):
				for j in range(nbly):
					for k in range(nblz):
					
							Kperm[i,j,k]=PetscP('',refDeg,kf[i*sx:(i+1)*sx,j*sy:(j+1)*sy,k*sz:(k+1)*sz],False,1e-4,0)

		
		tkperm[il]= time.time()-tkperm[il]
		np.save(datadir+'Kperm'+str(sx)+'.npy',Kperm)



	np.savetxt(rundir+'tkperm_sub.txt',tkperm)	

	return tkperm


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