#!/Users/weinberg.21/anaconda3/bin/python
"""
cosmodist2.py -- compute t_lb, comoving distance, angular diameter distance 
cosmodist2.py Omega_{m,0} Omega_{k,0} H0 w
  Omega_{m,0} = present day value of Omega_m
  Omega_{k,0} = present day value of Omega_k, set to zero for flat universe
  H0 = present day Hubble constant in km/s/Mpc
  w = dark energy (DE) equation of state parameter
  --- Omega_{DE,0} will be calculated as 1 - Omega_{m,0}-Omega_{k,0}-Omega{r,0}
Returns
  z
  Omega_m(z)
  t_loobkack(z)				(Gyr)
  d_c(z) = \int_0^z H_0/H(z') dz' 	(dimensionless)
  D_C(z) = (c/H_0) d_c(z)		(Mpc, for specified H0)
  D_M(z) = comoving ang. diam. distance	(Mpc)
  D_H(z) = c/H(z)			(Mpc)
  D_M/rd				(dimensionless)
  D_H/rd				(dimensionless)
  D_V/rd				(dimensionless)

The angle-averaged D_V = (z D_H D_M^2)^{1/3}  (eq. 2.6 of 2403.03002)
rd can be fixed or scaled from other parameters, as seen below.
"""

import numpy as np		# may be used in integrand
import sys			# needed for command line reads below
from scipy.integrate import quad	# quadrature integration

# functionality of a printf statement to stdout
def printf(format, *args):
    sys.stdout.write(format %args)

# read parameters
omega_m=float(sys.argv[1])
omega_k=float(sys.argv[2])
h0=float(sys.argv[3])
w=float(sys.argv[4])

# compute radiation density parameter omega_r, using values from Huterer tab 3.1
hh=h0/100.
omega_gamma=2.47e-5/hh**2
omega_nu=1.68e-5/hh**2
omega_r=omega_gamma+omega_nu

# compute omega_de
omega_de=1.-(omega_m+omega_k+omega_r)

# power-law scaling of rho_DE with (1+z)
wpow=3.0*(1.+w)

# compute sound horizon distance rd using eq. (2.5) of 2403.03002
#obh2=0.02218	# central value from eq. (2.8) of 2403.03002
#neff=3.04
#omh2=omega_m*(hh**2)
#rd=147.05*((omh2/0.1432)**-0.23)*((neff/3.04)**-0.1)*((obh2/0.02236)**-0.13)

# or fix to fiducial value from CMB constraints (quoted error is +/- 0.26)
rd=147.09

# comoving distance integrand
def dc_integrand(z):
    return(1./np.sqrt(omega_m*(1+z)**3+
                      omega_k*(1+z)**2+
		      omega_de*(1+z)**wpow+
		      omega_r*(1+z)**4))

# lookback time integrand
def tlookback_integrand(z):
    return(1./((1+z)*np.sqrt(omega_m*(1+z)**3+
                             omega_k*(1+z)**2+
			     omega_de*(1+z)**wpow+
			     omega_r*(1+z)**4)))

# time integrand forward from t=0, in a instead of z
def tforward_integrand(a):
    return(1./(a*np.sqrt(omega_m/(a**3)+
                         omega_k/(a**2)+
			 omega_de/(a**wpow)+
			 omega_r/(a**4) )))

dh0=2.997925e5/h0	# c/H0 in Mpc, if H0 is in km/s/Mpc
th=9.7778*100./h0	# 1/H0 in Gyr, if H0 is in km/s/Mpc

# compute present age of the universe
t0,err=quad(tforward_integrand,1.e-7,1.0)
t0 *= th

zout=np.array([0.30,0.51,0.71,0.93,1.32,1.49,2.33,1090])

printf("# H0=%.2f Om0=%.3f, Ode0=%.3f, Ok0=%.3f, w=%.3f, t0=%6.3f, rd=%6.2f\n",
          h0,omega_m,omega_de,omega_k,w,t0,rd)
printf("# z   O_m   tlook   d_c   D_C    D_M    D_H     D_M/rd  D_H/rd D_V/rd\n")

for z in zout:
    # comoving distance, in units of c/H0
    dc,err=quad(dc_integrand,0.0,z)	# quad returns integral and error

    # hubble distance, c/H(z), in Mpc
    dh=dh0*dc_integrand(z)

    # lookback time
    tlook,err=quad(tlookback_integrand,0.0,z)
    tlook *= th

    # angular diameter distance
    if (omega_k==0):
        da=dc*dh0/(1+z)
    elif (omega_k>0):
        da=dh0*np.sinh(np.sqrt(omega_k)*dc)/np.sqrt(omega_k)/(1+z)
    else:
        da=dh0*np.sin(np.sqrt(-omega_k)*dc)/np.sqrt(-omega_k)/(1+z)

    # convert to comoving angular diameter distance
    dm = da*(1+z)

    # value of Omega_matter 
    omega_mz=omega_m*(1+z)**3/(omega_m*(1+z)**3+
                               omega_k*(1+z)**2+
			       omega_de*(1+z)**wpow+
			       omega_r*(1+z)**4 )

    # scaled distances and dv (from eq. 2.6 of 2404.03002)
    dm1=dm/rd
    dh1=dh/rd
    dv1=(z*dh1*dm1**2)**(0.333333)

    printf("%4.2f %5.3f %6.3f %7.4f %6.1f %6.1f %7.2f %6.2f %7.3f %6.2f\n",
            z,omega_mz,tlook,dc,dc*dh0,dm,dh,dm1,dh1,dv1)