#! /usr/bin/env python

import os, sys
import numpy as np
from scipy import special
from scipy.optimize import fmin
#import scipy as sp

class Constants( object ):

        def __init__( self, Molecule ):

                ##############################
                # Some constants:
                self.kb      = 8.6173324E-5     # Boltzmann constant in eV/K
                self.cm2eV   = 1.23981E-4       # from cm-1 to eV
                self.J2eV    = 6.24150934E18    # from Joule to eV
                self.amu2Kg  = 1.660538921E-27  # from atomic mass unit (H = 1.008 amu) to Kg
                self.pi      = np.pi            # pi, imported from numpy

                if Molecule == 'D2':
                        ##############################
                        # D2 Constants:
                        self.gJeven   = 2.               # Ortho-para ratio constants 
                        self.gJodd    = 1.
                        self.Be       = 30.443   # Rovibrational constants (from NIST)
                        self.alphae   = 1.0786
                        self.omegae   = 3115.5
                        self.omegaxe  = 61.82
                        self.omegaye  = 0.562
                        self.De       = 0.01141
                        self.Betae    = 0.00022
                        self.MassAtom = 2.014
                        ##############################
                elif Molecule == 'H2':
                        ##############################
                        # H2 Constants:
                        self.gJeven   = 1.             # Ortho-para ratio constants 
                        self.gJodd    = 4.
                        self.Be       = 60.853         # Rovibrational constants (from NIST)
                        self.alphae   = 3.062
                        self.omegae   = 4401.21
                        self.omegaxe  = 121.33
                        self.omegaye  = 0.812
                        self.De       = 0.0471
                        self.Betae    = 0.0027
                        self.MassAtom = 1.008
                        ##############################

                self.MassMolecule = 2. * self.MassAtom


class VelocityDistribution( object ):

        def __init__( self, Vs, alpha, Molecule = 'D2' ):
        # Initialize distribution

                Const = Constants( Molecule  )

                self.Name = 'VelDistrib_Vs' + str( int(Vs) ) + '_alpha' + str( int(alpha) ) + '.dat'

                self.Es      = 1./2. * Const.MassMolecule * Const.amu2Kg * Vs ** 2.      * Const.J2eV
                self.DeltaEs = 2. * self.Es * ( alpha /  Vs )
                self.SqrtEs  = np.sqrt( self.Es )
                Srma    = Vs / alpha
                self.Normalization = ( self.DeltaEs / (2. * Srma) ) ** 2. * ( (1. + Srma **2.) * np.exp( -Srma**2. ) + np.sqrt( Const.pi ) * Srma * ( 1.5 + Srma**2. ) * (1. + special.erf( Srma) ) )
                self.Emax = (( self.DeltaEs**2. ) + 2. * ( self.Es**2. ) + 2. * self.Es * np.sqrt( ( self.Es**2. ) +  ( self.DeltaEs ** 2.)))/(4. * self.Es)
                print 'normal', self.Normalization

        def GetVelDistribution( self, E ):
        # Given a numpy vector "E" return a numpy vector containing f(E)dE

                # Define dE
                DeltaE = E[1] - E[0]

                Expon = 4. * self.Es * (( np.sqrt( E ) - self.SqrtEs ) / self.DeltaEs ) ** 2.
                gE    = E * np.exp( - Expon ) / self.Normalization

#               Eavg  =  sum ( gE * E * DeltaE ) 
#               VerifyNormalization = sum ( gE * DeltaE  )

                # Also generate a file containing distribution
                np.savetxt( self.Name, np.column_stack((E,gE)))

#               print ' The average collision energy is: (eV) ',     Eavg
#               print ' Verify normalization of the distribution: ', VerifyNormalization 

                return gE * DeltaE

        def GetVelAverage( self, E ):
                DeltaE = E[1] - E[0]
                Expon = 4. * self.Es * (( np.sqrt( E ) - self.SqrtEs ) / self.DeltaEs ) ** 2.
                gE    = E * np.exp( - Expon ) / self.Normalization

                Eavg  =  sum ( gE * E * DeltaE )

                VerifyNormalization = sum ( gE * DeltaE  )
                print ' Verify normalization of the distribution: ', VerifyNormalization

                return Eavg


if __name__ ==  '__main__':
 
  Vs    =5391.6  #2127.90
  alpha =1901.9  #297.967

  energies   = np.linspace( 0., 3., 3000 )
  VelDistrib = VelocityDistribution( Vs, alpha, Molecule='H2' )
  avgE       = VelDistrib.GetVelAverage( energies )
  print " Energy at maximum of the distribution: ", VelDistrib.Emax
  print " Average energy:                        ", avgE