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


def collect_scalar(filename):

	njobs = DotheLoop(-1)
	rdir='./data/'
	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 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

def searchError(filename):
	njobs = DotheLoop(-1)
	rdir='./data/'
	for job in range(njobs):	
		nclus=np.loadtxt(rdir+str(job)+'/'+filename)[:,4]
		for i in range(1,nclus.shape[0]):
			if nclus[0]!=nclus[i]:
				print(job,nclus[0],nclus[i])

	return

filename='resTestCon.txt'
searchError(filename)


res=collect_scalar(filename)
'''
stats = get_stats(res,0,True)
plot_keff(stats)
np.savetxt('Stats.txt',stats)
'''