#! /usr/bin/env python

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import scipy as sp
#from scipy import interpolate
from scipy.interpolate import griddata

Top_verticalFileName    = "2DCut_X4.503_Y0.000_theta0.00_phi0.00.dat"
BH_tiltedFileName       = "2DCut_X3.669_Y1.262_theta74.61_phi-122.27.dat"
Hollow_parallelFileName = "2DCut_X2.251_Y0.000_theta90.00_phi0.00.dat"

FileNames = [ Top_verticalFileName , BH_tiltedFileName , Hollow_parallelFileName ] 

mpl.rc("font", size=15 )

# Contour levels
levels = np.arange( -1.6, 1.0, 0.2 )

# Interpolated points in X, Y 
Xnew = np.linspace( 1.00, 3.00, 250 )
Ynew = np.linspace( 1.00, 3.75, 200 )

# Axis ranges 
Plt1Xmin = 1.0 ; Plt1Xmax = 1.6 ; Step1X = 0.2
Plt2Xmin = 1.0 ; Plt2Xmax = 2.2 ; Step2X = 0.4
Plt3Xmin = 1.0 ; Plt3Xmax = 3.0 ; Step3X = 0.5 

PltYmin = 1.0 ; PltYmax = 3.75

ticks = [ np.arange( Plt1Xmin, Plt1Xmax + Step1X, Step1X), \
          np.arange( Plt2Xmin, Plt2Xmax + Step2X, Step2X), \
          np.arange( Plt3Xmin, Plt3Xmax + Step3X, Step3X) ] 

fig, axarr = plt.subplots( 1, len( FileNames ), sharey=True )


for i in range( len( FileNames )):
	
	FileName = FileNames[i]
	Data = np.loadtxt( FileName )
	X, Y, Z = [], [], []

	for Row in Data:
		X.append( Row[0] )
		Y.append( Row[1] )
		Z.append( Row[2] )

	Xarray = np.array( X )
	Yarray = np.array( Y )
	Zarray = np.array( Z )

	Znew = griddata((Xarray, Yarray), Zarray, (Xnew[None,:], Ynew[:,None]), method='cubic')


	axarr[i].contourf( Xnew, Ynew, Znew, levels )#, colors='k' ) 
	axarr[i].contour( Xnew, Ynew, Znew, levels , colors='k' )
	#axarr[i].scatter(Xarray,Yarray,marker='o',c='b',s=5)
	axarr[i].set_xlim( [ ticks[i][0] , ticks[i][-1] ] ) 
	axarr[i].xaxis.set_ticks( ticks[i] )
#	axarr[i].set_xlim( [ ticks[i][0] , ticks[i][-1] - 0.1 ] )
#	print ticks[i]

#fig.subplots_adjust(wspace=0)
# Y label
#plt.subplot( 1, 3, 1)
axarr[0].set_ylabel( r'Z / $\AA$' ) 
axarr[0].set_ylim( [PltYmin, PltYmax ])

axarr[0].set_title( "Top-vertical" )
axarr[1].set_title( "Bridge/hollow-tilted" )
axarr[2].set_title( "Hollow-parallel" )
# X label
#axarr[1].subplot( 1, 3, 2)
axarr[1].set_xlabel( r'r / $\AA$' )

# Add position of the minima as +
axarr[0].scatter(1.141,2.694,marker='+',c='r',s=50)
axarr[1].scatter(1.316,1.544,marker='+',c='r',s=50)
axarr[2].scatter(1.370,1.391,marker='+',c='r',s=50)

# Add position of the barriers as x
axarr[0].scatter(1.12354,  3.67020, marker='x',c='k',s=30)
axarr[1].scatter(1.15213,  2.43709, marker='x',c='k',s=30)
axarr[2].scatter(1.15336,  2.41879, marker='x',c='k',s=30)

# Add coordinates
fig.text( 0.15, 0.20, r'$\theta = 0\degree$, $\phi = 0\degree$, ' " \n" '$X = 0$ $\AA$,' "\n" '$Y = 0$ $\AA$', size="large")
fig.text( 0.45, 0.55, r'$\theta = 74.6\degree$,' "\n" '$\phi = -122.3\degree$, ' " \n" '$X = 3.67$ $\AA$,' "\n" '$Y = 1.26$ $\AA$', size="large")
fig.text( 0.73, 0.55, r'$\theta = 90\degree$,  ' "\n" '$\phi = 0\degree$, ' " \n" '$X = 2.25$ $\AA$, ' "\n" '$Y = 0$ $\AA$', size="large")

# Add letters 
fig.text( 0.30, 0.85, "A" , size="x-large" )
fig.text( 0.58, 0.85, "B" , size="x-large" )
fig.text( 0.85, 0.85, "C" , size="x-large" )



plt.show()