#!/Users/weinberg.21/anaconda3/bin/python
"""
desi_chisq.py -- compute chisq of model relative to desi data
desi_chisq.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
  h0 omega_m omega_de omega_k w chisq1 chisq2 dm(1090)
  where the first 5 are just the run parameters
  chisq1 is a diagonal chi^2 relative to the DESI data
  chisq2 is a chi^2 that accounts for the anti-correlation of dm and dh errors
  dm(1090) is the comoving angular diameter distance to z=1090, divided by rd
  
"""

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
rd=147.05

# 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

# first do the redshifts with dm and dh measurements
zeff,dm_dat,dmerr,dh_dat,dherr,rho=np.loadtxt('desi.tbl1',usecols=(0,1,2,3,4,5),
                                       unpack=True)

# chisq1 is a diagonal chi^2, while chisq2 takes account of the anti-correlation
# between dm and dh errors

chisq1=0.0
chisq2=0.0
for i in range(len(zeff)):
    z=zeff[i]
    # 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)

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

    # scaled distances and dv (from eq. 2.6 of 2404.03002)
    dm1=dm/rd
    dh1=dh/rd
    dv1=(z*dh1*dm1**2)**(0.333333)
    dchisq = ((dm1-dm_dat[i])/dmerr[i])**2 + ((dh1-dh_dat[i])/dherr[i])**2
    chisq1 += dchisq
    dchisq =  (1./(1-(rho[i]**2))*(dchisq - 
              2.*rho[i]*(dm1-dm_dat[i])*(dh1-dh_dat[i])/(dmerr[i]*dherr[i])))
    chisq2 += dchisq

# add on the redshifts with dv measurements only
zeff,dv_dat,dverr=np.loadtxt('desi.tbl2',usecols=(0,1,2),unpack=True)

for i in range(len(zeff)):
    z=zeff[i]
    # 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)

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

    # scaled distances and dv (from eq. 2.6 of 2404.03002)
    dm1=dm/rd
    dh1=dh/rd
    dv1=(z*dh1*dm1**2)**(0.333333)
    dchisq = ((dv1-dv_dat[i])/dverr[i])**2 
    chisq1 += dchisq
    chisq2 += dchisq		

# compute dm(z=1090)/rd 
z=1090
dc,err=quad(dc_integrand,0.0,z)
if (omega_k==0):
    dm=dc*dh0
elif (omega_k>0):
    dm=dh0*np.sinh(np.sqrt(omega_k)*dc)/np.sqrt(omega_k)
else:
    dm=dh0*np.sin(np.sqrt(-omega_k)*dc)/np.sqrt(-omega_k)
dm1090=dm/rd

printf("%.2f %.3f %.3f %.4f %.3f  %.2f %.2f  %.4f\n",
       h0,omega_m,omega_de,omega_k,w,chisq1,chisq2,dm1090)