import yt
import numpy as np
import scipy.special as syspc
import scipy.integrate as syint



#
#	Load Simulation
#''''''''''''''''''''''''''''''
fn = "/home/pcguest/Documents/Simulations/enzo_tiny_cosmology/RD0009/RD0009"
print (fn)
ds = yt.load(fn) 		# load data
z= ds.current_redshift



#
#	Define few constants
#'''''''''''''''''''''''''''''''''
speed_of_light = ds.quan(2.997925e10, "cm/s")
boltzmann = ds.quan(1.380658e-16, 'erg/K')
m_proton = ds.quan(1.6726219e-24, 'g')
m_electron = ds.quan(9.10938e-28, 'g')
Y_helium = ds.quan(0.2, '')


#
#	Maxwell_Juttner Distribution (3-Dim)
#'''''''''''''''''''''''''''''''''''''''''''''''''''
def max_jut_dist(p, ne, Temperature) :
    func = ne / (4.0*np.pi * (m_electron.d * speed_of_light.d)**3)
    theta = boltzmann.d * Temperature.d / (m_electron.d * speed_of_light.d**2)
    gamma = np.sqrt(1.0 + (p / (m_electron.d * speed_of_light.d))**2)
    func /= theta
    func /= syspc.kn(2, 1.0/theta)
    func *= np.exp( -gamma/theta)
    return func


#
#	Piecewise integration
#''''''''''''''''''''''''''''''''''''
def int_fn (gm_min, gm_max, ne, Temperature) :
    gbreak = np.logspace(np.log10(gm_min), np.log10(gm_max), 50)
    sum1 = 0.0
    for g1, g2 in zip( gbreak[:-1], gbreak[1:] ) :
        Inte = syint.quad( max_jut_dist,  g1, g2, args=( ne, Temperature ) )
        sum1 += Inte[0]
    return sum1



#
#	Calculating Electron density
#''''''''''''''''''''''''''''''''''''''''''
def _edensity(field, data) :
    dens = data["density"] * (2.0 - Y_helium) / (2.0*m_proton)
    return dens
ds.add_field ("Edensity", function = _edensity, take_log = False, units = "cm**(-3)" )



#
#	Normalization contant
#''''''''''''''''''''''''''''''''''''''''''
def norm(field, data) :
    temp = int_fn(1.0e-34, 1.0e-18, data["Edensity", data["Temperature"]])
    return temp
ds.add_field("Norm", function = norm, take_log = False, units = "")