
"""PKTron Cosmology — Simple Universe Evolution Simulation."""
import numpy as np
_R = np.float32

def universe_step(state, dt=0.01, H0=67.4, Omega_m=0.315, Omega_L=0.685):
    """
    Evolve universe state by one time step using Friedmann equation.

    Parameters
    ----------
    state : dict with keys:
        a       : float  scale factor (a=1 today)
        t       : float  cosmic time in Gyr
        H       : float  Hubble parameter in km/s/Mpc
    dt    : float  time step in Gyr
    H0    : float  Hubble constant
    Omega_m, Omega_L : floats  matter and Lambda density params

    Returns
    -------
    new state dict
    """
    a = _R(state["a"])
    t = _R(state["t"])
    H = _R(state["H"])

    # Friedmann: H^2 = H0^2 * [Omega_m/a^3 + Omega_L]
    H_new = float(_R(H0)*np.sqrt(_R(Omega_m)/a**3+_R(Omega_L)))

    # a_dot = H*a  (in Gyr units: convert H from km/s/Mpc to Gyr^-1)
    H_Gyr  = _R(H) * _R(1.022e-3)   # km/s/Mpc -> Gyr^-1
    a_new  = float(a + H_Gyr*a*_R(dt))
    t_new  = float(t + dt)

    return {"a": a_new, "t": t_new, "H": H_new,
            "Omega_m_eff": float(_R(Omega_m)/a_new**3),
            "Omega_L_eff": float(Omega_L)}


def simulate_universe(n_steps=1000, dt=0.01, H0=67.4,
                      Omega_m=0.315, Omega_L=0.685, a0=0.01):
    """
    Simulate universe expansion from scale factor a0 to present.

    Returns
    -------
    dict: a_arr, t_arr, H_arr (all float32 arrays)
    """
    state = {"a": float(a0), "t": 0.0, "H": float(H0)}
    a_arr = np.zeros(n_steps, dtype=_R)
    t_arr = np.zeros(n_steps, dtype=_R)
    H_arr = np.zeros(n_steps, dtype=_R)

    for i in range(n_steps):
        a_arr[i] = state["a"]
        t_arr[i] = state["t"]
        H_arr[i] = state["H"]
        state     = universe_step(state, dt=dt, H0=H0,
                                  Omega_m=Omega_m, Omega_L=Omega_L)

    return {"a_arr": a_arr, "t_arr": t_arr, "H_arr": H_arr}


def big_bang_nucleosynthesis_fractions(eta=6.1e-10):
    """
    Approximate primordial element mass fractions from BBN.
    eta = baryon-to-photon ratio.
    Returns dict of mass fractions.
    """
    # Approximate analytic fits (Kolb & Turner 1990 inspired)
    Yp    = 0.2471 + 0.0016*(eta-6.1e-10)/1e-10   # He-4 mass fraction
    D_H   = 2.54e-5 * (6.1e-10/eta)**1.6           # Deuterium
    Li7_H = 4.7e-10                                 # Li-7
    return {"H_fraction" : round(1-Yp-D_H, 6),
            "He4_fraction": round(Yp, 6),
            "D_fraction"  : round(D_H, 8),
            "Li7_fraction": round(Li7_H, 12)}
