#!/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			#
#########################################################
#							#
# 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_cu111_h2.dat', 'energy_au111_hcl.dat' ]
titles 		= [ r'(a) H$_2$ + Cu(111)', '(b) HCl + Au(111)' ]  
references 	= [ 0.6920761661516792, 1.1103728608914025 ]

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

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

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

colormap = []
def ReadAndInterpolate(idx, INP, Title, reference, minE, maxE, spacing, bar=False):
	plt.subplot(Nrows,Ncols,idx)

	S = []
	THETA = []
	E = []

	lines = open(INP).readlines()
	for i in range(1, len(lines)):
		#if float(  lines[i].split()[0]  ) > -0.5 or float(  lines[i].split()[1]  ) > 90.: continue
		#if float(  lines[i].split()[1]  ) > 45.: continue
		#if float(  lines[i].split()[0]  ) > -0.5: continue
		if float(  lines[i].split()[2]  ) > 4.: continue
		S.append(float(  lines[i].split()[0]  ))
		THETA.append(float(  lines[i].split()[1]  ))
		E.append(float(  lines[i].split()[2]  )) 

	N = len(lines)
	S = np.array(S)
	THETA = np.array(THETA)
	E = ( np.array(E) - reference ) * 96.4853075
	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 = 100	# number of points per dimention in the denser grid
	Xnew, Ynew = np.meshgrid( np.linspace( min(S), max(S), density ),  np.linspace( min(THETA), max(THETA), density ) )

	# interpolate Z on the XY uniform grid
	#F 	= Rbf(S, THETA, E, function='cubic', smooth=0. )
	from scipy.interpolate import interp2d
	F = interp2d(S,THETA,E,kind='cubic')
	#Znew	= F(Xnew, Ynew)
	Znew = np.zeros( ( len( Xnew ), len( Ynew ) ) )
	for y in range( len(Ynew) ):
		for x in range( len(Xnew) ):
			Znew[x][y] = F( Xnew[x][y], Ynew[x][y] )

	areas 	 = plt.contourf(Xnew, Ynew, Znew, levels, zorder=-1, cmap='jet'  )# , colors='k' )	# coloured areas
	contours = plt.contour( Xnew, Ynew, Znew, levels, linewidths=1.0, colors='k', zorder=0 )	# black level lines
	colormap.append( areas )

	if bar: plt.colorbar(areas)
	
	plt.plot( [0.,0.], [0.,180.], color='k', ls='--', zorder=10 )
	
	if idx == 2:
		xaxis = reactioncoordinate( MEP_HCl[:,0], MEP_HCl[:,1], MEP_HCl[:,3] )
		plt.plot( xaxis, MEP_HCl[:,2], color='r' )
	if idx == 1:
		plt.plot( [-10., 10.], [90.,90.], color='r' )
	
	#plt.axis('scaled')
	plt.xlim(-1.2, 0.2)
	plt.ylim(min(THETA), max(THETA))
	if idx == 1:
		plt.yticks( [ 0., 45., 90., 135., 180. ] )
	if idx == 2:
		plt.xlabel( r'Reaction coordinate (\r{A})' )
		plt.yticks( [ 0., 45., 90., 135. ] )
	plt.ylabel( r'$\theta$ (degrees)' )

	plt.annotate(Title, xy=(0.05, 0.1), xycoords='axes fraction', bbox=dict(boxstyle='round', fc='w'))

	plt.tick_params(length=6, width=1, direction='in', top=True, right=True)

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

#Minimum Energy Path
# r [ Ang ]   Z [ Ang ]   E [ Ang ]
MEP_HCl = np.array(
[[1.2987, 2.9685, 113.74,  18.80],
[1.2994, 2.9370, 112.57,  20.15],
[1.3001, 2.9054, 111.36,  21.75],
[1.3007, 2.8736, 110.10,  23.42],
[1.3012, 2.8415, 108.83,  25.04],
[1.3020, 2.8091, 107.67,  26.67],
[1.3037, 2.7765, 106.85,  28.56],
[1.3055, 2.7435, 106.30,  30.69],
[1.3075, 2.7100, 105.85,  33.02],
[1.3104, 2.6762, 105.45,  35.67],
[1.3133, 2.6417, 104.98,  38.61],
[1.3168, 2.6068, 104.56,  41.86],
[1.3222, 2.5718, 104.45,  45.59],
[1.3283, 2.5365, 104.50,  49.79],
[1.3361, 2.5012, 104.78,  54.54],
[1.3471, 2.4671, 105.45,  59.96],
[1.3617, 2.4345, 106.60,  66.03],
[1.3838, 2.4069, 108.83,  72.70],
[1.4295, 2.3980, 113.54,  79.51],
[1.4927, 2.4065, 118.18,  85.49],
[1.5581, 2.4220, 121.06,  90.28],
[1.6335, 2.4529, 123.90,  93.69],
[1.6923, 2.4743, 125.32,  95.88],
[1.7373, 2.4861, 125.93,  97.35],
[1.7757, 2.4944, 126.20,  98.37],
[1.8084, 2.4987, 126.04,  99.09],
[1.8375, 2.5010, 125.54,  99.60],
[1.8641, 2.5019, 124.86,  99.97],
[1.8883, 2.5012, 124.12, 100.24],
[1.9110, 2.4993, 123.36, 100.43],
[1.9325, 2.4969, 122.64, 100.57],
[1.9533, 2.4942, 121.95, 100.65],
[1.9734, 2.4910, 121.28, 100.70],
[1.9929, 2.4873, 120.63, 100.71],
[2.0119, 2.4833, 119.99, 100.69],
[2.0307, 2.4788, 119.37, 100.64],
[2.0492, 2.4740, 118.76, 100.56],
[2.0676, 2.4689, 118.15, 100.46],
[2.0860, 2.4636, 117.54, 100.33],
[2.1044, 2.4579, 116.93, 100.18],
[2.1230, 2.4521, 116.32, 100.01],
[2.1418, 2.4461, 115.71,  99.81],
[2.1609, 2.4403, 115.11,  99.59],
[2.1803, 2.4350, 114.54,  99.34]]
)

def reactioncoordinate( r, Z, E ):
	if len(r) != len(Z) != len(E):
		print('lengths of r, Z and E are not equal')
		exit()
	
	s = [ 0. ]
	for i in range( 1, len( r ) ):
		dZ = Z[i] - Z[i-1]
		dr = r[i] - r[i-1]
		s.append( s[-1] + np.sqrt( dZ**2 + dr**2 ) )
	ds = s[np.argmax(E)]
	for i in range( len(s) ):
		s[i] -= ds
	
	return s

for i in range( len( inputs ) ):
	ReadAndInterpolate(i+1, inputs[i], titles[i], references[i], minE, maxE, spacing, False)

plt.tight_layout()
plt.subplots_adjust(wspace=0.0, hspace=0.0)

cbar_ax = fig.add_axes([0.82, 0.13, 0.025, 0.82]) #xmin, ymin, xwidth, ywidth
fig.colorbar(colormap[0], cax=cbar_ax) #, ticks=np.arange(0.0,0.7,0.1)
cbar_ax.set_ylabel('Energy (kJ/mol)')
fig.subplots_adjust(right=0.8)

plt.savefig('TS_elbow_theta.pdf')
#plt.show()