#!/usr/bin/env python

######################################################### 
import 	matplotlib.pyplot as plt			#	
import 	matplotlib        as mpl			#
import 	numpy             as np				#
from  	scipy.interpolate import Rbf 			#
from	scipy.optimize    import minimize		#	
from   	math              import *			#
from   	sys		  import stdout			#
from   	phi_scan	  import distance, psi_interp1D	#
#########################################################
#							#
# this code produces a 2D contour plot and the MEP on 	#
# the PES. 						#
# this code is written for a diatomic molecule thinking	#
# about the 2D cut of the potential considering the	#
# interatomic distance (r) and the distance of the 	#
# CoM (Z) from the surface as variables. 		#
# Therefore the data.dat input file must be formatted   # 
# as:		r      Z      E				#
# You can define the variables related to the plot	# 
# apperance and (more important) to the interpolation   #
# and the MEP calculation in the VARIABLES section. 	# 
#							# 	
# Davide Migliorini		      Leiden, 2015	# 				      	
#							#
#-------------------------------------------------------#
# A massive update has been implemented including a way #
# better fitting procedure to find the MEP and a better #
# 2D interpolation function: scipy.interpolate.Rbf	#
# I am too lazy to proper document all the changes	#
# however the comments in the code should make it 	#
# (fairly) self-explainatory.				#
# I hope.						#
#							#
# Davide Migliorini			Leiden, 2018	#
#							#
#-------------------------------------------------------#
# Some changes were implemented to improve the figures, #
# too lazy to document.					#
#							#
# Nick Gerrits				Leiden, 2018	#
#							#
#########################################################

#########################################################
# VARIABLES #############################################

# Title &  file name
inputs 		= [ 'energy_NN_constraint.dat' ]
titles 		= [ 'Cu(111) - NN - Elbow plot' ]  
references 	= [ -23.135431420346208 * 96.4853075 ]

# contour line
minE    	=  -20.00			# min and max E shown in the plot
maxE    	=  190.			#	
spacing 	=  10.0		# E spacing between contour lines

# set up plot parameters
mpl.rc("text", usetex=True)
fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(3.69,5))

########################################################
# Contour plot #########################################
########################################################

def ReadAndInterpolate(INP, Title, reference, minE, maxE, spacing, plot,bar=False):

	Z = []
	r = []
	E = []

	lines = open(INP).readlines()
	for i in range(1, len(lines)):
		Z.append(float(  lines[i].split()[0]  ))
		r.append(float(  lines[i].split()[1]  ))
		E.append(float(  lines[i].split()[2]  )) 

	N = len(lines)
	r = np.array(r)
	Z = np.array(Z)
	E = np.array(E) * 96.4853075 -reference
	print " -----------------------"
	print " file: ", INP
	print " title:", Title
	print " Energy values (min-max):"
	print " (", min(E), max(E), ")  [ eV ]"

	# CONTOUR LINES PLOT ###################################

	levels = np.arange( minE, maxE , spacing )    #range of energies shown

	# generate a uniform XY grid 
	density = 200	# number of points per dimention in the denser grid
	Xnew, Ynew = np.meshgrid( np.linspace( min(r), max(r), density ),  np.linspace( min(Z), max(Z), density ) )

	# interpolate Z on the XY uniform grid
	F 	= Rbf(r, Z, E, function='linear', smooth=0.5 )
	Znew	= F(Xnew, Ynew)		

	# Contour lines ###
#	plt.subplot(1,2,plot)

#	areas 	 = plt.contourf(Xnew, Ynew, Znew, levels, zorder=-1  )# , colors='k' )	# coloured areas
	contours = plt.contour( Xnew, Ynew, Znew, levels, colors='k', zorder=0 )	# black level lines
	ax.clabel(contours, contours.levels, inline=True, fmt='%3.0f')
#	plt.scatter(Xnew,Ynew, s=1)

	# compute Minimum Energy Path
	MEP 	= psi_interp1D(r,Z,E,F)
	MEP_r	= [ x[0] for x in MEP ] 
	MEP_Z	= [ x[1] for x in MEP ] 
	MEP_E	= [ x[2] for x in MEP ] 
	pos_max = np.argmax(MEP_E)

#	plt.plot(MEP_r, MEP_Z, '--', linewidth='4', color = 'white')
#	plt.scatter(MEP_r, MEP_Z, color = 'black', zorder=1)
	plt.scatter(MEP_r[pos_max], MEP_Z[pos_max], color = 'black', marker='s', s=45, zorder=2)

	print "Minimum Energy Path"
	print " r [ Ang ]   Z [ Ang ]   E [ Ang ]"
	for j in range(len(MEP)):
		print " %10.6f  %10.6f  %10.6f " % tuple(MEP[j])

	# plot appereance ###
#	plt.colorbar(areas)
	
	#plt.axis('scaled')
	plt.xlim(min(r), max(r))
	plt.ylim(min(Z), max(Z))
	plt.ylabel( r'Z (\r{A})' )
	plt.xlabel( r'r (\r{A})' )
#	plt.title( "%s" % Title )

###################################################################
###################################################################
###################################################################

def plot_std(filename, label):
	plt.plot([2.25, 2.25], [0.,1.], linestyle='--', color='k')
	plt.xlim(1.0, 2.7)
	plt.ylim(0,0.34)
	plt.ylabel('Distribution (arb.)')
	plt.xlabel('Z of carbon ($\mathrm{\AA}$)')
	plt.tick_params(length=6, width=1, labelleft=False)

	Zc = []
	Zcdat = open(filename, 'r').readlines()
	for i in range(len(Zcdat)):
		Zc.append( float(Zcdat[i]) )

	nZc, bins = np.histogram( Zc, bins=10, density=True )
	xdata = []
	xdata.append( bins[0] - (bins[1] - bins[0]) / 2. )
	for i in range(len(bins)-1):
		xdata.append( ( bins[i] + bins[i+1] ) / 2. )
	xdata.append( (bins[-1] - bins[2]) / 2. + bins[-1] )
	nZc = np.insert(nZc, 0, 0)
	nZc = np.append(nZc, 0)

	plt.plot( xdata, nZc/sum(nZc), label=label )
#	plt.hist( Zc, bins=20, normed=True, histtype='step' )

	plt.legend(loc='best', numpoints=1, borderaxespad=0.2, handletextpad=0.5, frameon=False)
	plt.annotate('(a)', xy=(2.5, 0.3), size=12)

plt.subplot(2,1,1)

plot_std('../../../zc/summaries/zc_181v0.dat', 'Ground state')	#0.00030, 0.00009
plot_std('../../../zc/summaries/zc_143v1.dat', r'$v_1=1$')	#0.00015, 0.00004
plot_std('../../../zc/summaries/zc_97v2.dat', r'$v_1=2$')	#0.00040, 0.00014

ax = plt.subplot(2,1,2)

ReadAndInterpolate(inputs[0], titles[0], references[0], minE, maxE, spacing, 1, True)

def plot_traj(filename, color, label):
	Z = []
	r = []
	data = open(str(filename), 'r').readlines()
	for i in range( len(data) ):
		Z.append( float(data[i].split()[1]) )
		r.append( float(data[i].split()[2]) )
	plt.plot( r, Z, linewidth=1, c=color, label=label )

plot_traj('181v0_112348.dat', 'b', 'Ground state')
plot_traj('143v1_014889.dat', 'g', r'$\nu_1=1$')
plot_traj('97v2_014829.dat', 'r', r'$\nu_1=2$')

#plt.legend(borderaxespad=0.)
plt.annotate('(b)', xy=(1.92, 2.9), size=12)

plt.tight_layout()
plt.subplots_adjust(hspace = 0.28)
plt.savefig('elbow_zc_traj.pdf')
#plt.show()