import numpy as np
from scipy.special import erf, erfinv
from scipy.optimize import curve_fit
from scipy.stats import norm
from FFTMA import gen
from tools.generation.config import DotheLoop, get_config	
from tools.generation.variograma import get_lc
from scipy.interpolate import interp1d
import sys
import time
import os
#from memory_profiler import profile

#@profile
def fftmaGenerator(datadir,job,conffile):

	t0=time.time()
	parser, iterables = get_config(conffile)	
	params = DotheLoop(job,parser, iterables )

	binary=parser.get('Generation','binary')
	uselc_bin=parser.get('Generation','lcBin')

	if binary=='yes':
		logn='no'
		con,lc,p,seed = params[0],params[1],params[2],params[3]
		variance=0
	else:
		logn='yes'
		con,lc,variance,seed = params[0],params[1],params[2],params[3]
		p=0


	Nx,Ny,Nz = int(parser.get('Generation','Nx')), int(parser.get('Generation','Ny')), int(parser.get('Generation','Nz'))
	#N=int(42.666666667*lc)
	#Nx,Ny,Nz = N,N,N
	#print(N)

	kh,kl,vario = float(parser.get('Generation','kh')),float(parser.get('Generation','kl')), int(parser.get('Generation','variogram_type'))
	compute_lc=parser.get('Generation','compute_lc')
	generate_K(Nx,Ny,Nz,con,lc,p,kh,kl,seed,logn,variance,vario,datadir,compute_lc,uselc_bin)

	np.savetxt(datadir+'GenParams.txt',np.array([time.time()-t0,Nx,Ny,Nz,con,lc,p,kh,kl,seed,variance,vario]),header='Runtime,Nx,Ny,Nz,con,lc,p,kh,kl,seed,variance,vario')


	return

	
def obtainLctobin(p,con,vario):


	lc=np.load('./tools/generation/lc.npy',allow_pickle=True, encoding = 'latin1').item()

	f=interp1d(lc['p'],lc[vario,con])
	if p==0 or p==1:
		return 1.0
	return 1.0/f(p)
	
def obtainLctobinBack(p,con,vario):

	pb=np.linspace(0.0,1.0,11)


	if vario==2:
		i=[0.0, 1.951, 2.142, 2.247, 2.301, 2.317, 2.301, 2.246, 2.142, 1.952, 0.0]
		c=[0.0, 1.188, 1.460, 1.730, 2.017, 2.284, 2.497, 2.652, 2.736, 2.689, 0.0]
		d=[0.0,2.689, 2.736,2.652, 2.497, 2.284, 2.017, 1.730, 1.460, 1.188, 0.0]
		lcBin=np.array([i,c,d])
		lcBin=lcBin/3.0

	if vario==1:
		i=[0.0,3.13, 3.66, 3.94, 4.08, 4.10, 4.01, 3.84, 3.55, 3.00,0.0]
		c=[0.0,0.85, 1.095, 1.312, 1.547, 1.762, 1.966, 2.149, 2.257, 2.186,0.0]
		d=[0.0,2.186, 2.2575,2.1495,1.9660,1.7625,1.5476,1.3128,1.0950,0.8510,0.0]
		lcBin=np.array([i,c,d])
		lcBin=lcBin/6.0

	f=interp1d(pb,lcBin[con-1])
	return 1.0/f(p)

 
#@profile
def generate_K(Nx,Ny,Nz,con,lc,p,kh,kl,seed,logn,LogVariance,vario,datadir,compute_lc,uselc_bin):


	k=genGaussK(Nx,Ny,Nz,con,lc,p,kh,kl,seed,logn,LogVariance,vario,uselc_bin)
	if compute_lc =='yes':
		lcG=get_lc(k,vario)
		lcNst=lcG
		lcBin=np.nan
	if con==2:
		k  = -nst(k)  #normal score transform
		if compute_lc =='yes':
			lcNst=get_lc(k,vario)
	if con==3:
		k = nst(k)
		if compute_lc =='yes':
			lcNst=get_lc(k,vario)

	if logn == 'yes':
		k=k*(LogVariance**0.5)
		k = np.exp(k)
	
	else:
		k = binarize(k,kh,kl,p)
		if compute_lc =='yes':
			lcBin=get_lc(np.where(k>kl,1,0),vario)
	np.save(datadir+'k.npy',k)
	if compute_lc =='yes':
		np.savetxt(datadir+'lc.txt',np.array([lcG,lcNst,lcBin]),header='lcG, lcNst, lcBin')
	return

def genGaussK(Nx,Ny,Nz,con,lc,p,kh,kl,seed,logn,LogVariance,vario,uselc_bin):


	typ=0  #structure du champ: 0=normal; 1=lognormal; 2=log-10
	dx, dy, dz = 1.0, 1.0, 1.0
	var=1	#Nbr de structure du variogramme
	alpha=1 #valeur exposant
	if con==0:
		lc=0.000001

	if (con==2 or con==3) and vario==2:
		lc=lc/0.60019978939
	if (con==2 or con==3) and vario==1:
		lc=lc/0.38165155120015

	if uselc_bin=='yes' and con!=0:
		lc=lc*obtainLctobin(p,con,vario)
	v1 = (var, vario, alpha, lc, lc, lc, 1, 0, 0, 0, 1, 0)  # coord des vecteurs de base (1 0 0) y (0 1 0)
	k=gen(Nz, Ny, Nx, dx, dy, dz, seed, [v1], 0, 1, 0)  # 0, 1, 0 = mean, variance, typ	  #Generation of a correlated standard dsitribution N(0,1)

	return k


def nst(kc):
	kc=np.abs(kc)
	kc=np.sqrt(2)*erfinv(2*erf(kc/np.sqrt(2))-1)
	return kc


def binarize(kc,kh,kl,p):
	
	if kc.size < 100**3:
		if p>0:
			at=int((1-p)*kc.size)  #
		else:
			at=kc.size-1
		t1=np.sort(kc.reshape(-1))[at] #get permeability treshold
		kc=np.where(kc<t1, kl,kh)    #Binarization
		t1=0
	else:
		t1=norm.ppf(1-p)
		kc=np.where(kc<t1, kl,kh)
	return kc

#CONFIG_FILE_PATH = 'config.ini' if 'CONFIG_FILE_PATH' not in os.environ else os.environ['CONFIG_FILE_PATH']

#fftmaGenerator(sys.argv[1],int(sys.argv[2]),CONFIG_FILE_PATH)