#! /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.485_Y0.000_theta0.00_phi0.00.dat"
BH_tiltedFileName       = "2DCut_X3.657_Y1.264_theta74.48_phi-121.89.dat"
Hollow_parallelFileName = "2DCut_X2.243_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.672, 1.537, 1.378 ]

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)

	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()