import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, LinearSegmentedColormap

def plotK(kk,pdir,logn):

	y=np.arange(kk.shape[0])
	x=np.arange(kk.shape[1])
	newcolors = np.zeros((2,4))
	alto = np.array([0.0, 0.0, 0.0, 1])
	bajo = np.array([191/256.0, 191/256.0, 191/256.0, 1]) #[108.0/256, 122.0/256, 137.0/256, 1])

	alto = np.array([204.0/254, 0.0, 0.0, 1])
	bajo = np.array([0.0, 0.0, 153.0/254, 1]) #[108.0/256, 122.0/256, 137.0/256, 1])
	newcolors[0, :] = bajo
	newcolors[1, :] = alto
	newcmp = ListedColormap(newcolors)

	
	if logn==True:
		kk=np.log(kk)
		vmin,vmax=-2*np.var(kk)+np.mean(kk),2*np.var(kk)+np.mean(kk)
		#print(vmax)
		colormap='viridis'
		plt.pcolormesh(x,y,kk,cmap=colormap)#,vmin=vmin,vmax=vmax)
	else:
		#colormap='binary'
		plt.pcolormesh(x,y,kk,cmap=newcmp)

	cbar=plt.colorbar()
	cbar.set_label('k')
	#plt.title('Guassian N(0,1)')
	plt.savefig(pdir+'k.png')
	plt.close()
	'''
	if logn==True:
		plt.hist(kk.reshape(-1),range=(2*vmin,2*vmax),histtype='step',bins=250,density=True)
		plt.xlabel('k')
		plt.ylabel('p(k)')
		plt.savefig(pdir+'histo.png')
	'''
	
	return



def plotK_imshow(kk,pdir,logn):
	kk=np.rot90(kk)
	y=np.arange(kk.shape[0])
	x=np.arange(kk.shape[1])
	newcolors = np.zeros((2,4))
	alto = np.array([0.0, 0.0, 0.0, 1])
	bajo = np.array([191/256.0, 191/256.0, 191/256.0, 1]) #[108.0/256, 122.0/256, 137.0/256, 1])

	alto = np.array([204.0/254, 0.0, 0.0, 1])
	bajo = np.array([0.0, 0.0, 153.0/254, 1]) #[108.0/256, 122.0/256, 137.0/256, 1])
	newcolors[0, :] = bajo
	newcolors[1, :] = alto
	newcmp = ListedColormap(newcolors)

	
	if logn==True:
		kk=np.log(kk)
		vmin,vmax=-3*np.var(kk)+np.mean(kk),3*np.var(kk)+np.mean(kk)
		#print(vmax)
		colormap='viridis'
		plt.imshow(kk,vmin=vmin,vmax=vmax) #,cmap='binary'
	else:
		#colormap='binary'
		plt.imshow(kk,cmap='binary') #,cmap='binary'

	plt.colorbar()
	#cbar.set_label('k')
	#plt.title('Guassian N(0,1)')
	plt.tight_layout()
	plt.savefig(pdir+'k.png')
	plt.close()
	'''
	if logn==True:
		plt.hist(kk.reshape(-1),range=(2*vmin,2*vmax),histtype='step',bins=250,density=True)
		plt.xlabel('k')
		plt.ylabel('p(k)')
		plt.savefig(pdir+'histo.png')
	'''
		
	return

def plot_hist(k,pdir,logn):

	plt.figure(1)
	if logn==True:
		k=np.log(k)
		vmin,vmax=-4*np.var(k)+np.mean(k),4*np.var(k)+np.mean(k)
		plt.hist(k.reshape(-1),range=(vmin,vmax))
	else:
		plt.hist(k.reshape(-1))
	plt.xlabel('k')
	plt.ylabel('Counts')
	plt.savefig(pdir+'-histo.png')
	plt.close()
	return



rdir='./perco_lc8/'

for i in range(11):

	
	k=np.load(rdir+str(i)+'/k.npy')[:,:,0]
	log='False'

	plotK_imshow(k,rdir+str(i)+'Map',log)
	#plot_hist(k,rdir+'Res/'+resname,log)