#!/bin/env python

################################################################
# For Hubble constant and Omega_m values, see
#
# Bennett, C.L. et al The 1% Concordance Hubble Constant Astrophysical Journal 794 (2014)
#
# Martin Krzywinski
# v0.1
# 4 Jul 2018

import sys
import argparse
import math

parser = argparse.ArgumentParser(description="Calculate light-travel and comoving distances for objects with a given redshift. Set your own cosmological parametesr or use the reasonable defaults provided (see Bennett, C.L. et al The 1% Concordance Hubble Constant Astrophysical Journal 794 (2014)). The output is the current age of the universe, the age of the object when it emitted the light, the light-travel distance to the object and the comoving distance to the object. Units are either Gly (giga light-years) or Gy (giga years).",epilog="by Martin Krzywinski, http://mkweb.bcgsc.ca/unvierse-superclusters-and-voids",formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-z','--z',   type=float, default=0.2, help='redshift')
parser.add_argument('-Wr','--Wr', type=float, default=8.59798189985467e-5, help='relativistic mass density')
parser.add_argument('-Wk','--Wk', type=float, default=0, help='curvature')
parser.add_argument('-Wm','--Wm', type=float, default=0.286, help='mass density')
parser.add_argument('-WV','--WV', type=float, help='dark matter density')
parser.add_argument('-n','--n',   type=float, default=100000,help='integration steps')

args = parser.parse_args()

if not args.WV:
  args.WV = 1 - args.Wr - args.Wk - args.Wm

# Hubble parameter, as function of redshift
def Ez(z,Wr,Wm,Wk,WV):
    zz = (1+z)
    return(math.sqrt(Wr*zz**4 + Wm*zz**3 + Wk*zz**2 + WV))

# Hubble parameter, as function of a = 1/(1+z)
def Ea(a,Wr,Wm,Wk,WV):
    return(math.sqrt(Wr/a**2 + Wm/a + Wk + WV*a**2))

H0   = 69.6            # Hubble constant
c    = 299792.458      # speed of light, km/s
pc   = 3.26156         # parsec to light-year conversion
mult = (c/H0)*pc/1e3   # integrals are in units of c/H0, this factor brings us to Gy or Gly
sum_comoving = 0
sum_light    = 0
sum_univage  = 0
sum_univsize = 0
z = args.z
a = 1/(1+z)

# int(0..1) dx/Ea(x)  = age of universe
# int(a..1) dx/Ea(x)  = light-travel distance
# int(a..1) dx/xEa(x) = comoving distance

n = args.n
for i in range(n):
    f   = (i+0.5)/n
    x   = a + (1-a) * f # a .. 1
    xx  = f             # 0 .. 1
    ex  = Ea(x,args.Wr,args.Wm,args.Wk,args.WV)
    exx = Ea(xx,args.Wr,args.Wm,args.Wk,args.WV)
    sum_comoving += (1-a)/(x*ex)
    sum_light    += (1-a)/(  ex)
    sum_univsize += 1/(xx*exx)
    sum_univage  += 1/(   exx)

results = [mult/n*i for i in [sum_univage,sum_univsize,sum_univage-sum_light,sum_light,sum_comoving]]
print("z {:.2f} U {:.3f} Gy {:.3f} Gly T0 {:.3f} Gy T {:.3f} Gly C {:.3f} Gly".format(args.z,*results))