#cosmic age for cosmology with w0-wa dark energy and massive neutrino
#by Zhiqi Huang, huangzhq25@mail.sysu.edu.cn 
#use c/H_0 as the length unit, and 1/H_0 the time unit, but finally convert to Gyr for output

from scipy.integrate import quad
import numpy as np
import matplotlib.pyplot as plt


z = 0. #where you want to compute the age


#dark energy equation of state = w0 + wa (1-a)
w0 = -1.
wa = 0.
#cold dark matter + baryon; this does not include massive neutrinos
Omega_m = 0.3
#spatial curvature
Omega_k = 0.
#effective number of species 
N_eff = 3.046 
#masses of species of neutrinos in eV
mass_nu = np.array([0.001, 0.009, 0.05])
#Hubble constant in km/s/Mpc
H0 = 70. 
#CMB temperature
T_CMB = 2.726


## integrals for massive neutrinos
def I_rho(lam):
    cutdown = 0.5
    cutup = 7.
    k = 1./6.9
    if lam < cutdown:
        return (7*np.pi**2/240.)*(1.+ (5./7./np.pi**2/k)*lam**2)**k
    elif lam > cutup:
        return 9.13453712e-02*lam + 5.90977955e-01/lam -4.52/(lam+0.86)**3
    else:
        lnlam = np.log(lam)
        return  0.30627596287 + lnlam*(0.033956 + lnlam*(0.02977417  + lnlam*(0.01602 + lnlam*(0.005584 + lnlam*1.15e-3))))

def I_p(lam):
    cutdown = 0.5
    cutup = 7.
    k = -1./14.
    if lam <= cutdown:
        return 7*np.pi**2/720.*(1.- (5./7./np.pi**2/k)*lam**2)**k
    elif lam > cutup:
        return 3.93985303e-01/lam -6.03424688/lam**3 + 2.54842801e2/(lam+1.15)**5
    else:
        lnlam = np.log(lam)
        return 0.09077334191+lnlam*(-0.00853080347+lnlam*(-0.0060953041+lnlam*( -0.00210631549374+lnlam*( -4.3851383e-05 + lnlam * 3.87e-4))))


    
#derived parameters
#photon density parameter
Omega_gamma = 2.4748e-5/(H0/100.)**2*(T_CMB/2.726)**4
#neutrino density parameter
eV_in_Kelvin = 11604.505
T_CNB = T_CMB*(4./11.)**(1./3.)
Omega_nu_reference = 0.56202e-5/(H0/100.)**2 * N_eff/len(mass_nu)/I_rho(0.) # this is, Omega_nu/I_rho(0) assuming massless
Omega_nu = np.array([I_rho(m*eV_in_Kelvin/T_CNB) * Omega_nu_reference for m in mass_nu ])
#dark energy density parameter
Omega_L = 1.0 - Omega_m - Omega_k - Omega_gamma - sum(Omega_nu)

## rho_DE/rho_{DE0} as a function of a
def rhoDE_ratio(a):
    return np.exp(-3.*((1.+w0+wa)*np.log(a)+wa*(1.-a)))

##rho_nu/rho_critical  as a function of a
def Omega_nu_scaled(a):
    s = 0.    
    for i in range(len(mass_nu)):
        s +=  I_rho(mass_nu[i]*eV_in_Kelvin*a/T_CNB)
    return s/a**4*Omega_nu_reference

#Hubble parameter H = dot a / a as a function of a, in unit of H0
def Hofa(a):
    return np.sqrt(Omega_L*rhoDE_ratio(a) + Omega_m/a**3 + Omega_k /a**2 + Omega_gamma/a**4 + Omega_nu_scaled(a))

## dt / da in unit of H0^{-1}
def dtbyda(a):
    return 1.0/a/Hofa(a)


age, err = quad(dtbyda, 1.e-20, 1./(1.+z))  #age is in unit 1/H0, I start from 1.e-20 instead of 0 to avoid possible log(0) and 1/0 errors
#now convert to Gyr
Gyr = 1.e9*(365.2422*24*3600)  #Gyr in SI unit
H0Gyr = H0*(1.e3/3.086e22)*Gyr #H0 in unit of Gyr^{-1}
print("At z = ", z, ", Age = " + str(age/H0Gyr) + " Gyr") #convert to Gyr