#! /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
from scipy.optimize import fmin


class Interpolation:

	def __init__( self, X, Y, Z ):
		self.f = interpolate.interp2d( X, Y, Z, kind='cubic') 
		return

	def Value( self, Xnew, Ynew):
		return self.f( Xnew, Ynew )

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 )

rEquil = 1.1172392

ZMin = [ 2.694, 1.544, 1.391 ]

OutputFileNames = [ "TopVertical-MEP.dat", "BridgeHollowTilted-MEP.dat", "HollowParallel-MEP.dat"] 
 
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')

	V = Interpolation(Xarray, Yarray, Zarray)

#	for value in np.linspace( 1.00, 1.5, 50 ):
#		print value, V.Value( value, 1.378 )[0]
#	quit()

	File = open( OutputFileNames[i], "w" )

	rGuess = 1.2

	for ZValue in np.linspace( 4.00, ZMin[i], 100 ):
		rValue = fmin( V.Value, rGuess , args = tuple([ ZValue ]) )[0]
		string = str( " % 12.5f % 12.5f % 12.5f \n" %tuple( [ ZValue, V.Value( rValue, ZValue )[0], rValue ] ) )
		File.write( string )

	del V
	File.close()
#	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
#axarr[0].set_ylabel( r'Z / $\AA$' ) 
#axarr[0].set_ylim( [PltYmin, PltYmax ])

# X label
#axarr[1].set_xlabel( r'r / $\AA$' )

# Titles
#axarr[0].set_title( "Top-vertical" )
#axarr[1].set_title( "Bridge/hollow-tilted" )
#axarr[2].set_title( "Hollow-parallel" )

# Add position of the minima as +
#axarr[0].scatter(1.137,2.672,marker='+',c='r',s=50)
#axarr[1].scatter(1.307,1.537,marker='+',c='r',s=50)
#axarr[2].scatter(1.363,1.378,marker='+',c='r',s=50)

# 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.5\degree$,' "\n" '$\phi = -121.9\degree$, ' " \n" '$X = 3.66$ $\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.24$ $\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()