#luminosity distance and angular diameter distances 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 Mpc for output

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


z = 2. #where you want to compute the distance

#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.315
#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 = 67.4 
#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 z
def rhoDE_ratio(z):
    return np.exp(3.*((1.+w0+wa)*np.log(1.+z)-wa*z/(1+z)))

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

##pressure/density of neutrinos as a function of z
def w_nu(z):
    rho = 0.
    p = 0.
    for i in range(len(mass_nu)):
        lam = mass_nu[i]*eV_in_Kelvin/(1.+z)/T_CNB
        rho +=  I_rho(lam)
        p += I_p(lam)
    return p/rho

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

def dchibydz(z): 
    return 1.0/Hofz(z)

def rofchi(chi):
    if(abs(Omega_k)<1.e-20):
        return chi
    if(Omega_k < 0):
        return np.sin(np.sqrt(-Omega_k)*chi)/np.sqrt(-Omega_k)
    return np.sinh(np.sqrt(Omega_k)*chi)/np.sqrt(Omega_k)

def ComovingAngularDiameterDistance(z): #r
    chi, err = quad(dchibydz, 0., z)
    r = rofchi(chi)
    return r

def AngularDiameterDistance(z): #d_A
    return ComovingAngularDiameterDistance(z)/(1.+z)

def LuminosityDistance(z): #d_L
    return ComovingAngularDiameterDistance(z)*(1.+z)

r = ComovingAngularDiameterDistance(z)


#convert to Mpc
speed_of_light = 2.998e5 ##speed of light in km/s
H0Mpc = H0/speed_of_light # H0*Mpc/c 
print("At z = " + str(z))
print("luminosity distance  = "+str(r*(1+z)/H0Mpc)+" Mpc")
print("angular diameter distance  = "+str(r/(1+z)/H0Mpc)+" Mpc")