# 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 cosmological look-back time from redshift in Python:

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

# 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
plt.plot([x for x in rs if x<fo], [y for x,y in zip(rs,lb) if x<fo], color='red')
plt.plot([x for x in rs if x>fo], [y for x,y in zip(rs,lb) if x>fo], color='darkred')
plt.text(13.2, 9.5, 'Furthest observation as of 2023:\n' +
     'the metal-poor JADES-GS-z13-0 galaxy\nat z=13.2: 13.4 Gyr ago', ha='center')

plt.title('Look-back Time by Redshift')
plt.xlabel('z: (observed λ - expected λ) / expected λ')
plt.ylabel('Billion Years Ago')
plt.xticks(range(21))
plt.yticks(list(range(14)) + [age0])
plt.text(-0.5, 13.78, "Big Bang", va='center')
plt.gca().yaxis.set_major_formatter(ticker.FormatStrFormatter('%.1f'))
plt.grid(True, color='lightgray')
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['top'].set_visible(False)

for t in range(0, 13):
    z = rs[min(range(len(lb)), key=lambda i: abs(lb[i]-t))]
    plt.text(z, t, f"   z = {z:.2f}", ha='left', va='center', fontsize='small')

for z in range(7, 20, 2):
    t = zt(z)
    plt.text(z, t - 0.2, f"{t:.2f}", ha='center', va='top', fontsize='small')

for z in range(6, 21, 2):
    t = zt(z)
    plt.text(z, t + 0.1, f"{t:.2f}", ha='center', va='bottom', fontsize='small')

plt.savefig('time_by_redshift.png', bbox_inches='tight')
#plt.show()  # https://i.ibb.co/LpdYXNx/time-by-redshift.png