#luminosity distance and angular diameter distances for w0wa model
#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
#if you want the result to be in h^{-1}Mpc, set H0 = 100 (i.e. h=1)

import numpy as np
from scipy.integrate import quad

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

#density parameters
Omega_m = 0.3
Omega_k = 0.
Omega_r = 0.0 #for z < 100 ignore the radiation
Omega_L = 1.0 - Omega_m - Omega_k - Omega_r
w0 = -0.8
wa = 0.2

H0 = 70  #km/s/Mpc, it is used only when you want the results in Mpc


## 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)))


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

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_L
    return ComovingAngularDiameterDistance(z)/(1.+z)

def LuminosityDistance(z): #d_A
    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")