from astropy.cosmology import Planck18, z_at_value, FlatLambdaCDM
import astropy.units as u
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from astropy.constants import k_B
from scipy.integrate import quad

%matplotlib widget

t = """Particles	Type\tMass	spin	color	g\tQCD
γ 	boson\t     0 	2 	1 	2\tFalse
W+,W−	boson\t     80.4GeV	3 	1 	6\tFalse
Z	boson\t     91.2GeV	3 	1 	3\tFalse
Gluons 	boson\t     0 	2 	8 	16\tTrue
Higgs 	boson\t     125.3GeV	1 	1 	1\tFalse
u,u¯	fermion\t     2MeV	2 	3 	12\tTrue
d,d¯	fermion\t    5MeV	2 	3 	12\tTrue
s,s¯	fermion\t    100MeV	2 	3 	12\tTrue
c,c¯	fermion\t    1GeV	2 	3 	12\tTrue
b,b¯	fermion\t    4GeV	2 	3 	12\tTrue
t,t¯	fermion\t    173GeV	2 	3 	12\tTrue
e+,e−	fermion\t    0.511MeV	2 	1 	4\tFalse
μ+,μ−	fermion\t    105.7MeV	2 	1 	4\tFalse
τ+,τ−	fermion\t    1.777GeV	2 	1 	4\tFalse
pions	boson\t    135MeV	1 	1 	3\tFalse
νe,νe¯	fermion\t    <3eV	1 	1 	2\tFalse
νμ,νμ¯	fermion\t    <0.19MeV	1 	1 	2\tFalse
ντ,ντ¯	fermion\t    <18.2MeV	1 	1 	2\tFalse
"""

with open("particle_g.csv", 'w') as f:
    f.write(t)

df = pd.read_csv("particle_g.csv", delimiter="\t")

df

mass_txt = df["Mass"]
mass_MeV = []
for m in mass_txt:
    if "eV" in m:
        if "<" in m:
            mass_MeV.append(0)
        elif "GeV" in m:
            mass_MeV.append(float(m.split('GeV')[0])*1e3)
        elif "MeV" in m:
            mass_MeV.append(float(m.split('MeV')[0]))
    else:
        mass_MeV.append(float(m))
mass_MeV = np.array(mass_MeV)

mass_T = mass_MeV / 1.38e-23 * 1.6e-19 * 1e6

a = 2.725 / mass_T
z = 1/a - 1
z[z < 0] = np.inf

T0 = Planck18.Tcmb0  # K
cosmo = FlatLambdaCDM(H0=Planck18.H0, Om0=Planck18.Om0, Tcmb0=T0, m_nu=0)
Omegar0 = cosmo.Ogamma0 + cosmo.Onu0  # with massless neutrinos
Omegar0

invH0 = (1/Planck18.H0).to(u.s).value
invH0

def inv_aH(a):
    return (invH0/(a*(np.sqrt(Omegar0 / a**4))))

def a_to_time(a):
    return quad(inv_aH, 0, a)[0]

def T_to_time(T):
    return a_to_time(T0.to(u.K).value / T)

def T_to_MeV(T):
    return T * 1.38e-23 / (1.6e-19 * 1e6)

def MeV_to_time(MeV):
    T = MeV / 1.38e-23 * 1.6e-19 * 1e6
    return T_to_time(T)

MeV_to_time(150)

t = [MeV_to_time(mass_MeV[k]) for k in range(len(mass_MeV))]

df["MeV"] = mass_MeV
df["a"] = a
df["z"] = z
df["t"] = t
df.sort_values("MeV")

def gs_tot(MeV, entropy=False):
    if MeV > 150:
        gb = np.sum(df["g"][(df["MeV"] < MeV) & (df["Type"]=="boson") & (df["Particles"]!="pions")])
        gf = np.sum(df["g"][(df["MeV"] < MeV) & (df["Type"]=="fermion")])
    elif MeV < 150 and MeV >= 0.511:  # QCD phase transition
        gb = np.sum(df["g"][(df["MeV"] < MeV) & (df["Type"]=="boson") & (df["QCD"]==False)])
        gf = np.sum(df["g"][(df["MeV"] < MeV) & (df["Type"]=="fermion") & (df["QCD"]==False)])
    elif MeV < 0.511:  # e+/e- annihilation
        gb = 2
        if not entropy:
            gf = 2 * 3 * (4/11)**(4/3)  # neutrinos at lower T
        else:
            gf = 2 * 3 * (4/11)
    return gb + 7/8 * gf

gs_tot(1e6)

MeVs = np.sort(np.concatenate([np.logspace(-2, 7, 1000), df["MeV"]]))
ts = [MeV_to_time(MeV) for MeV in MeVs]
gs = [gs_tot(MeV) for MeV in MeVs]

fig = plt.figure()
plt.plot(MeVs, gs)
plt.grid()
plt.xscale("log")
plt.xlim((1e6, 1e-1))
ax1 = plt.gca()
#ax1.invert_xaxis()
ax2 = plt.gca().twiny()
ax2.set_xscale("log")
ticks = ax1.get_xticks()[::-1][1:-1]
ticks = ticks[(ticks>=MeVs.min()) & (ticks<=MeVs.max())]
#print(ticks)
#ticks = 10.**np.arange(6, -3, -1)
#print(ticks)
new_ticks = [MeV_to_time(tick) for tick in ticks]
#print(new_ticks)
ax2.set_xticks(new_ticks)
ax2.set_xticklabels([f"{tick:.0g}" for tick in new_ticks])
ax2.set_xlabel("$t$ [s]", fontsize=16)
ax1.set_ylabel("$g_*(T)$", fontsize=16)
ax1.set_xlabel("$T$ [MeV]", fontsize=16)
fig.tight_layout()
plt.show()

def T_precise(a):
    T_tmp = T0 / a
    gs0 = gs_tot(T_to_MeV(T0).value, entropy=True)
    for k in range(10):
        T_tmp = T0 * (gs0/gs_tot(T_to_MeV(T_tmp).value, entropy=True))**(1/3) / a
    return T_tmp

T_precise(a=1e-5)

a = np.logspace(-14, -8, 100)
Ts = [T_precise(aa).value for aa in a]

fig = plt.figure()
ax1 = plt.gca()
ax1.plot(a, [T_to_MeV(t) for t in Ts], label="$T(a, g_{\star S})$")
ax1.plot(a, [T_to_MeV(T0.value/aa) for aa in a], 'k--', label=r"$T_0\times (a_0/a)$")
ax1.grid()
ax1.set_xscale("log")
ax1.set_yscale("log")
ax1.set_xlim(a.min(), a.max())
ax1.legend(fontsize=14)
#ax1.invert_xaxis()
ax2 = plt.gca().twiny()
ax2.set_xscale("log")
ticks = ax1.get_xticks()[::-1]
ticks = ticks[(ticks>=a.min()) & (ticks<=a.max())]
#print(ticks)
#ticks = 10.**np.arange(6, -3, -1)
#print(ticks)
new_ticks = [a_to_time(tick) for tick in ticks]
#print(new_ticks)
ax2.set_xticks(new_ticks)
ax2.set_xticklabels([f"{tick:.2g}" for tick in new_ticks])
ax1.set_xlabel("$a/a_0$", fontsize=16)
ax1.set_ylabel("$T$ [MeV]", fontsize=16)
ax2.set_xlabel("$t$ [s]", fontsize=16)
fig.tight_layout()
plt.show()

T_to_MeV(1.8e10)