# Thanks to ChatGPT-4 and the Fortran-90 code from arxiv:1303.5961,
#     https://code.google.com/archive/p/cosmonom/downloads
# here's how to get the age of the universe from redshift in Python:

from scipy.special import hyp2f1  # hypergeometric function 2F1 is in integral solution
import matplotlib.pyplot as plt

# Cosmological parameters from the Fortran params.f90 header
#H0 = 67.15       # Hubble constant in km/s/Mpc (or, 73.5: the "crisis in cosmology")
H0 = 69.32        # from Explainxkcd for 2853: Redshift; seems a consensus compromise
#OL = 0.683       # Cosmological constant for dark energy density, Omega_Lambda or _vac
#Om = 0.317       # Density parameter for matter, Omega_mass
Om = 0.286        # From https://arxiv.org/pdf/1406.1718.pdf page 8
OL = 1.0 - Om - 0.4165/(H0**2)  # flat curvature, from https://www.astro.ucla.edu/~wright/CC.python
                  # (on https://www.astro.ucla.edu/~wright/CosmoCalc.html which see)
#print(f"{OL=:.3F}")  # 0.714

# Age of universe at redshift z as a closed-form solution to its integral definition,
def age_at_z(z):  # ...which is 27 times faster than the original numeric integration
    hypergeom = hyp2f1(0.5, 0.5, 1.5, -OL / (Om * (z + 1)**3))
    return (2/3) * hypergeom / (Om**0.5 * (z + 1)**1.5) * (977.8 / H0)  # 977.8 for Gyr

# Current age of the universe at redshift 0 in Gyr
age0 = age_at_z(0)  # 13.78

# Function to calculate the look-back time at redshift z in Gyr
def zt(z):  # from the function name in the Fortran cosmonom.f90 code
    return age0 - age_at_z(z)

rs = [z * 20 / 299 for z in range(300)]  # redshifts 0 to 20 in 300 steps
lb = [zt(z) for z in rs]  # look_back_times

fo = 13.2  # furthest observation at present
#print(age_at_z(fo))  # 0.3285

rs = [z * 15 / 299 + 5 for z in range(300)]  # redshifts 5 to 20 in 300 steps
ages = [age_at_z(z) * 1000 for z in rs]  # Gyr to million years

plt.plot([x for x in rs if x<fo], [y for x,y in zip(rs,ages) if x<fo], color='red')
plt.plot([x for x in rs if x>fo], [y for x,y in zip(rs,ages) if x>fo], color='darkred')
plt.text(13.2, 650, 'Furthest observation as of 2023:\n' +
        'the metal-poor JADES-GS-z13-0 galaxy\nat z=13.2: age 329 Myr', ha='center')

plt.title('Age of Universe by Redshift')
plt.xlabel('z: (observed λ - expected λ) / expected λ')
plt.ylabel('Million Years')
plt.xticks(range(5, 21))
plt.yticks(range(0, 1300, 100))
plt.grid(True, color='lightgray')
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['top'].set_visible(False)

plt.savefig('age_by_redshift.png', bbox_inches='tight')
#plt.show()  # https://i.ibb.co/C537rxJ/age-by-redshift.png