#!/usr/bin/python

import numpy as np    
from scipy import interpolate 
from scipy import linalg as LA 
from scipy import integrate 
from random import uniform 

# ==============================================================================
# Universal constants  
# ==============================================================================

# for more information: http://docs.scipy.org/doc/scipy/reference/constants.html
from scipy.constants import physical_constants as phc
atmas = phc['atomic mass constant'][0]
eV2J  = phc['electron volt-joule relationship'][0]  
Bc    = phc['Boltzmann constant'][0]
hbar  = phc['Planck constant'][0] / (2*np.pi)
lsp   = phc['speed of light in vacuum'][0]
H2J   = phc['Hartree energy'][0]

# ==============================================================================
# INPUT ------------------------------------------------------------------------        
# ==============================================================================

# load input  
data=np.genfromtxt('Input/input.txt', comments="#", skip_header=0, invalid_raise=False, loose=True, dtype=float)

# required atomic properties
mass1 = data[0] * atmas                 # atomic mass units 
mass2 = data[1] * atmas                 # atomic mass units

Tn   = data[2]                          # nozzle temperature
M    = int(data[3])                     # number of grids for the finite difference
Nvs  = int(data[4])                     # maximum number of vibrational states 
Jmax = int(data[5])                     # maximum angular momentum quantum number
et   = data[6]                          # Threshold energy for the harmonic fit in eV 

# load interaction potential  
data=np.loadtxt('Input/pot.txt',dtype=float)
r=data[:,0]
V=data[:,1]

# ==============================================================================
# Define required functions     
# ==============================================================================

# clear screen
import os
def cls():
    os.system('cls' if os.name=='nt' else 'clear')

cls()

# Morse Potential Fit
def Morse(x,a):
    global Re, Ed
    return Ed*( 1.0-np.exp(-a*(x-Re)) )**2.0

# Harmonic potential
def quadratic(x,a):
    global Re
    return 0.5*a*(x-Re)**2 

# Boltzmann distribution for rovibrational states
def Fb(v,J,Tn):
    global Erv
    bta=Bc*Tn
    return (2*J+1)*np.exp(-Erv[v,J]/bta )

# ==============================================================================
# spline interpolation
# ==============================================================================

# Calculate grid points. The point r_min with V(r_min)=min(V) is also involved.  
re=r[V==min(V)][0]
dr=(max(r)-min(r))/M
rnew=[re]
while min(r)+dr<=rnew[len(rnew)-1]:
      rnew.append(rnew[len(rnew)-1]-dr)

rnew.append(re+dr)
while rnew[len(rnew)-1]<=max(r)-dr:
      rnew.append(rnew[len(rnew)-1]+dr)

rnew=np.sort(rnew) 

from scipy import interpolate as si
tck = si.splrep(r, V, s=0)
Vnew = si.splev(rnew, tck, der=0) 

V=V-min(Vnew)
Vnew=Vnew-min(Vnew)

# Fitting the interaction energy with a Morse-like potential
Re=rnew[Vnew==min(Vnew)][0]        # equilibrium position with minimum energy 
Ed=Vnew[len(Vnew)-1]-min(Vnew)     # dissociation energy in eV

Redm=mass1*mass2/(mass1+mass2)     # reduced mass


# ==============================================================================
# Hamiltonian - diagonalization
# ==============================================================================

# calculate "we" using the following relation: E(neu)/h.c = we (neu+1/2) - we*xe (neu+1/2)**2 + we*ye (neu+1/2)**3 + ... + cte
powl=5    # power of the last neu-dependent term in above expansion 

# diagonalization
coeff=-hbar**2/(2.0*Redm)
dr=(rnew[1]-rnew[0])*1e-10
Erv=np.zeros((max(Nvs+1,powl),Jmax+1))
for J in range(0,Jmax+1):
    H = np.diag(-coeff*2/dr**2 + eV2J*Vnew - coeff*J*(J+1)/(rnew*1.0e-10)**2, 0) + \
        np.diag(coeff/dr**2 * np.ones((len(Vnew)-1)), 1) + np.diag(coeff/dr**2 * np.ones((len(Vnew)-1)), -1)
    E, sai = LA.eigh(H)
    for v in range(0, Nvs+1):
         Erv[v,J]=E[v]

Xmat=[]
Emat=[]
for i in range(0,powl):
    xm=[]
    limit=range(1,powl)     
    limit.append(0)
    #for j in limit:   # to ignore "cte" in the expansion, replace the for-loop limits by """for j in range(1,powl+1):"""
    for j in range(1,powl+1):   
        xm.append((i+1./2.)**j)
    Xmat.append(xm)
    Emat.append([Erv[i,0]/(hbar*2*np.pi*lsp*100)])

wmat=np.dot(LA.inv(Xmat),Emat)

# calculate frequency   
w0=2*np.pi*lsp*100*wmat[0][0]

# ==============================================================================
# general infromation of the rotational states  
# ==============================================================================
Be=hbar/(4.0*np.pi*lsp*100.*Redm*(Re*1e-10)**2)

# ==============================================================================
# Output files
# ==============================================================================

# write Erv(v,J) down on Erov.txt
g=open('Output/Erov.txt', 'w')
for v in range(0, Nvs):
    for J in range(0,Jmax+1):
        g.write( "%5.8f " %tuple(np.array([ Erv[v,J]/eV2J ])) )
    g.write( "\n" )

g.close()

g=open('Output/Intpot.txt', 'w')
for i in range(0, len(Vnew)):
    g.write( "%5.10f %5.10f\n" %tuple(np.array([ rnew[i],Vnew[i] ])) )

g.close()

g=open('Output/Syspar.txt', 'w')
g.write( "%5.8E      # w (H)\n" %tuple(np.array([ w0 ])) )
g.write( "%5.8E      # we (cm-1)\n" %tuple(np.array([ wmat[0][0] ])) )
g.write( "%5.8E      # we*xe (cm-1)\n" %tuple(np.array([ abs(wmat[1][0]) ])) )
g.write( "%5.8E      # we*ye (cm-1)\n" %tuple(np.array([ wmat[2][0] ])) )
g.write( "%5.8E      # Be (cm-1)\n" %tuple(np.array([ Be ])) )
g.close()

print "\n"
print "================================================================================"
print "Rovibrational eigenenergies E(v,J) have been stored, in eV, on Outputs/Erov.txt.\nYou can also find it below on the screen." 
print "Moreover, given below is some general information regarding the characteristics\nof the vibrational and rotational states."
print "\n"
print "================================================================================"
print "******** Structural information about the vibrational states ********\n"
print "\t1) we    = ", wmat[0][0], "(1/cm)"
print "\t2) we*xe = ", abs(wmat[1][0]), "(1/cm)"
print "\t3) we*ye = ", wmat[2][0], "(1/cm)\n"
print "--------------------------------------------------------------------------------"
print "Estimation of the frequency and period of vibrations ---------------------------\n"
print "\n\n\tf =",w0/(2*np.pi)/1e12, "(TH) and T =",2*np.pi*1e12/w0, "(ps)" 

print "\n"
print "================================================================================"
print "******** Parametres of the rotational states ********\n"
print "\t Be = hbar/4.pi.c.I =", Be, "(1/cm)"

print "\n"
print "================================================================================"
print "          Zero point energy and the lowest excited states "
print "            including the rotational degree of freedom    \n"
for J in range(0,Jmax+1):
    print " J=", J
    print "\t   E (eV)         E (Hartree)    E/h.c (cm-1)      DeltaE (eV)"
    for v in range(0,Nvs):
        pr = [v, Erv[v,J]/eV2J, Erv[v,J]/H2J, Erv[v,J]/(hbar*2*np.pi*lsp*100), (Erv[v+1,J]-Erv[v,J])/eV2J]     
        print( "%5s    %12.8f \t %12.8f \t %12.6f \t %12.8f" %tuple(pr) )
    print "\n"
 
print "Boltzmann distribution at given Tn for differenet rotational \nand vibrational quantum numbers:\n"
for v in range(0, Nvs):
    mylist=[]
    for J in range(0,Jmax+1):
         mylist.append( Fb(v,J,Tn)/(2*J+1) )
    np.set_printoptions(precision=4)
    print(np.array(mylist))