#!/usr/bin/python3.6
from numpy import *

from math import *
import sys

'''
program vibpartition_eV.py
	- determines harmonic vibrational contribution to enthalpy (eV), entropy (eV/K) and Gibbs 
	free energy (eV) using quantum mechanical expressions

	- note that, within harmonic approximation, enthalpy is exactly equivalent to internal 
	energy and Gibbs energy is equivalent Helmholtz energy
	 	
	- example of usage:

	python3 ../../tools/thermo/vibpartition_eV.py 440. refdata/freq.dat	

	here the temperature (the first argument) is set to 440 K and harmonic frequencies are 
	read from the file refdata/freq.dat created by vibfreq.py 
'''

class Thermodynamics:
    def __init__(self, sigma, T, p, nfile='vib', mysmall=1.):
        AVOGADRO = 6.02214e23
        self.PLANCK = 6.62608e-34
        self.BOLTZMANN = 1.38065e-23
        self.CL = 2.997925e08
        self.eVTOJ = 1.60218e-19

        qvib, Uvib, Hvib, Svib, zpe = self.part_vib(T, nfile, mysmall)
        # print 'qvib:',qvib

        Uvib /= self.eVTOJ
        Hvib /= self.eVTOJ
        Svib /= self.eVTOJ
        zpe /= self.eVTOJ

        print(qvib)
        print('')
        # print 'T [K]','Hvib [J/mol]','ZPE [J/mol]','Hvib-ZPE [J/mol]'
        # print T,Hvib,zpe,Hvib-zpe
        # print 'T [K]','Hvib [J/mol]','Svib [J/K/mol]','Gvib [J/mol]'
        # print T,Hvib,Svib,Hvib-T*Svib
        print('T [K]', 'ZPE [eV]', 'Hvib-ZPE [eV]', 'Hvib [eV]', 'Svib [eV/K]', 'Gvib [eV]')
        print(T, zpe, Hvib-zpe, Hvib, Svib, Hvib-T*Svib)

    def part_vib(self, T, nfile, mysmall):
        EXP = 2.7182818
        modes = open(nfile)
        vibs = []
        vib_count = 0
        vib_total = 0
        for line in modes.readlines():
            # vib_total+=1
            row = line.split()
            # read in wavenumbers in cm-1
            if len(row) == 11:
                if row[1] == "f" and row[2] == "=":
                    vib_total += 1
                    vm = float(row[3])
                    if vm > mysmall:
                        vib_count += 1
                        # convert to energy in J
                        vm = vm*self.PLANCK*self.CL*1e2
                        vibs.append(vm)
                elif row[1] == "f/i":
                    vib_total += 1
        modes.close()
        print("nDOF:", vib_total)
        print("nVIB:", vib_count)
        zpe = 0.5*sum(vibs)
        # print "ZPE (J):", zpe/1.60218e-19

        qvib = 1
        Uvib = 0.
        Hvib = 0.
        Svib = 0.
        if T > 0.:
            for i in range(len(vibs)):
                part = 1-EXP**(-vibs[i]/(self.BOLTZMANN*T))
                part2 = EXP**(vibs[i]/(self.BOLTZMANN*T))-1.
                Uvib += vibs[i]/part2
                Hvib += vibs[i]*(0.5+1./(part2))
                Svib += (vibs[i]/(T*part2)-self.BOLTZMANN*log(part))
                part = EXP**(-vibs[i]/(2*self.BOLTZMANN*T))/part
                qvib = qvib*part
        else:
            Hvib = zpe
        return qvib, Uvib, Hvib, Svib, zpe


mysmall = 1.
T = float(sys.argv[1])
nfile = sys.argv[2]
print(nfile)
Thermodynamics(2, T, 1e5, nfile, mysmall)
