#!/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	#
from   	phi_scan_edit_top	import psi_interp1D_edit_top	#
#########################################################
#							#
# 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.dat' ]
titles 		= [ 'a (E)', r'b ($\theta$)' ]  
references 	= [ 1715.14915092 * 96.4853075 ]

# contour line
minE    	=  40.			# min and max E shown in the plot
maxE    	=  170.			#	
spacing 	= 10.0		# E spacing between contour lines

# set up plot parameters
mpl.rc("text", usetex=True)
#plt.rcParams.update({'font.size': 12})
Nrows=1
Ncols=2
fig, ax = plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(6.69,3.))

########################################################
# Contour plot #########################################
########################################################
images = []
def ReadAndInterpolate(INP, Title, reference, minE, maxE, spacing, plot,bar=False):
	plt.subplot(Nrows,Ncols,plot)

	Z = []
	r = []
	THETA = []
	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]  ))
		THETA.append(float(  lines[i].split()[2]  ))
		E.append(float(  lines[i].split()[4]  )) 

	N = len(lines)
	r = np.array(r)
	Z = np.array(Z)
	THETA = np.array(THETA)
	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='cubic', smooth=0. )
	FT	= Rbf(r, Z, THETA, function='cubic', smooth=0. )
	if plot == 1:
		Znew	= F(Xnew, Ynew)
	else:
		Znew	= FT(Xnew, Ynew)

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

	if plot == 1:
		areas 	 = plt.contourf(Xnew, Ynew, Znew, levels, zorder=-1, cmap='jet'  )# , colors='k' )	# coloured areas
	else:
		areas 	 = plt.contourf(Xnew, Ynew, Znew, levels, zorder=-1, cmap='viridis'  )# , colors='k' )	# coloured areas
	contours = plt.contour( Xnew, Ynew, Znew, levels, colors='k', zorder=0 )	# black level lines
#	plt.scatter(Xnew,Ynew, s=1)

	Zmep = 2.59
	rmep = 1.88
	dz = 0.01
	dr = 0.01
	dt = 1e-4
	MEP_Z = [ Zmep ]
	MEP_r = [ rmep ]
	for j in range( 200 ):
		Zp = ( F( rmep, Zmep + dz/2. ) - F( rmep, Zmep - dz/2. ) ) / dz
		rp = ( F( rmep + dr/2., Zmep ) - F( rmep - dr/2., Zmep ) ) / dr
		Zmep += -Zp * dt
		rmep += -rp * dt
		MEP_Z.append( Zmep )
		MEP_r.append( rmep )

	Zmep = 2.58
	rmep = 1.90
	dz = 0.01
	dr = 0.01
	dt = 1e-4
	for j in range( 200 ):
		Zp = ( F( rmep, Zmep + dz/2. ) - F( rmep, Zmep - dz/2. ) ) / dz
		rp = ( F( rmep + dr/2., Zmep ) - F( rmep - dr/2., Zmep ) ) / dr
		Zmep += -Zp * dt
		rmep += -rp * dt
		MEP_Z.append( Zmep )
		MEP_r.append( rmep )

	# compute Minimum Energy Path
	MEP2 	= psi_interp1D_edit_top(r,Z,E,F)
	MEP2	= np.delete( MEP2, np.argmax( MEP2[:,2] ), 0 )
	MEP2_r	= [ x[0] for x in MEP2 ] 
	MEP2_Z	= [ x[1] for x in MEP2 ] 
	MEP2_E	= [ x[2] for x in MEP2 ] 
	pos_max = np.argmax(MEP2_E)

#	plt.plot(MEP_r, MEP_Z, '--', linewidth='4', color = 'white')
	plt.scatter(MEP_r, MEP_Z, color = 'white', zorder=1)
	plt.scatter(MEP2_r, MEP2_Z, color = 'grey', zorder=1)
	plt.scatter(MEP2_r[pos_max], MEP2_Z[pos_max], color = 'black', marker='s', s=35, 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]) )

#	bbox = dict(boxstyle="round", fc="w")
	plt.annotate(Title, xy=(2., 2.85), size=14)

	# plot appereance ###
	if bar:
		plt.colorbar(areas)
	images.append( areas )
	
	#plt.axis('scaled')
	plt.xlim(min(r), max(r)-0.01)
	plt.ylim(min(Z)+0.01, max(Z))
	if plot == 1:
		plt.ylabel( r'$Z_{\rm Cl}$ (\r{A})' )
	plt.xlabel( r'$r$ (\r{A})' )
#	plt.title( "%s" % Title )

	plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
	if plot == 2:
		plt.tick_params(labelleft=False)

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

ReadAndInterpolate(inputs[0], titles[0], references[0], 0., 210., 10., 1, False)
ReadAndInterpolate(inputs[0], titles[1], references[0], 40., 170., 10., 2, False)

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

cbar_ax = fig.add_axes([0.09, 0.1, 0.025, 0.85]) #xmin, ymin, xwidth, ywidth
fig.colorbar(images[0], cax=cbar_ax) #, ticks=np.arange(0.0,0.7,0.1)
cbar_ax.set_ylabel('Energy (kJ/mol)')
cbar_ax.tick_params(axis='y',labelleft=True, labelright=False, right=False, left=True)
cbar_ax.yaxis.set_label_position("left")

cbar_ax2 = fig.add_axes([0.89, 0.1, 0.025, 0.85]) #xmin, ymin, xwidth, ywidth
fig.colorbar(images[1], cax=cbar_ax2) #, ticks=np.arange(0.0,0.7,0.1)
cbar_ax2.set_ylabel(r'$\theta$ (degrees)')
fig.subplots_adjust(right=0.87, left=0.19)

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