import numpy as np
import matplotlib.pyplot as plt
from tools.generation.config import DotheLoop, get_config
import os

def collect_scalar(filename,rdir):

	njobs = DotheLoop(-1)
	res=np.array([])
	for job in range(njobs):	
		res=np.append(res,np.loadtxt(rdir+str(job)+'/'+filename))

	res=res.reshape(njobs,-1)
	return res

def get_stats(res,col,logv):

	parser, iterables = get_config()

	seeds=iterables['seeds']
	n_of_seeds=len(seeds)
	ps = iterables['ps']
	n_of_ps=len(ps)
	stats=np.zeros((n_of_ps,3))
	x=res[:,col]
	if logv==True:
		x=np.log(x)

	for i in range(n_of_ps):

		stats[i,0]=ps[i]
		stats[i,1]=np.nanmean(x[i*n_of_seeds:(i+1)*n_of_seeds])
		stats[i,2]=np.nanvar(x[i*n_of_seeds:(i+1)*n_of_seeds])


	if logv==True:
		stats[:,1]=np.exp(stats[:,1])

	return stats

def collect_Conec(scales,rdir):

	parser, iterables = get_config(rdir+'config.ini')

	ps = iterables['ps'] 
	njobs = DotheLoop(-1,parser, iterables)
	res=dict()
	for job in range(njobs):	
		for scale in scales:
			try:
				fdir=rdir+str(job)+'/ConnectivityMetrics/'+str(scale)+'.npy'
				jobres=np.load(fdir).item()

				params=DotheLoop(job,parser,iterables)
				indp=int(np.where(ps == params[2])[0])
				for ckey in jobres.keys():
					try:
						res[params[0],params[1],scale,ckey,indp]=np.append(res[params[0],params[1],scale,ckey,indp],jobres[ckey].reshape(-1))
					except KeyError:
						res[params[0],params[1],scale,ckey,indp]=jobres[ckey].reshape(-1)			
			except  IOError:
				pass
	return res


def ConValidat(conkey,scale,ddir):

	scales=[scale]
	resdict=collect_Conec(scales,ddir)
	parser, iterables = get_config(ddir+'config.ini')
	params=DotheLoop(0,parser,iterables)
	con=params[0]
	lc=params[1]
	x,y,yv=constasP(con,lc,scales[0],conkey,resdict,iterables)

	try:
		os.makedirs('./plots/'+ddir)
	except:
		pass
	plt.figure(1)
	plt.plot(x,y,marker='x')
	plt.xlabel('p')
	plt.ylabel(conkey)
	plt.grid()
	plt.savefig('./plots/'+ddir+conkey+str(scale)+'.png')
	plt.close()

	return

def showValidateResults(conkeys):
	for conkey in conkeys:
		ConValidat(conkey,128,'./data_Val2D/')
		ConValidat(conkey,16,'./data_Val3D/')

	return
def constasP(con,lc,scale,conkey,res,iterables):


	x = iterables['ps'] 

	y = np.zeros((x.shape))
	vy = np.zeros((x.shape))
	for i in range((x.shape[0])):

		y[i]=np.mean(res[con,lc,scale,conkey,i])
		vy[i]=np.mean(res[con,lc,scale,conkey,i])

	return x,y,vy


def plot_keff(stats):

	ylabel=r'$K_{eff}$'
	xlabel=r'$p$'
	fsize=14
	plt.figure(1)
	plt.semilogy(stats[:,0],stats[:,1])
	plt.xlabel(xlabel,fontsize=fsize)
	plt.ylabel(ylabel,fontsize=fsize)
	plt.grid()
	plt.savefig('Keff_p.png')
	plt.close()

	plt.figure(2)
	plt.plot(stats[:,0],stats[:,2])
	plt.xlabel(xlabel,fontsize=fsize)
	plt.ylabel(ylabel,fontsize=fsize)
	plt.grid()
	plt.savefig('vKeff_p.png')
	plt.close()	
	return


showValidateResults(['P','S','npx','Plen'])