
"""
PKTron Cosmology — Hubble Expansion & Cosmic Distance Ladder
=============================================================
Friedmann equations, Hubble law, co-moving distances.
Constants are SI unless noted.
"""
import numpy as np
_R = np.float32

# Physical constants
c_km_s    = _R(299792.458)   # km/s
H0_planck = _R(67.4)          # km/s/Mpc  (Planck 2018)
Mpc_m     = _R(3.0857e22)     # metres per Megaparsec


def hubble_expansion(distance_Mpc, H0=None):
    """
    Recession velocity from Hubble law: v = H0 * d

    Parameters
    ----------
    distance_Mpc : float  co-moving distance in Mpc
    H0           : float  Hubble constant km/s/Mpc (default 67.4)

    Returns
    -------
    dict: velocity_km_s, distance_Mpc, H0
    """
    if H0 is None: H0 = float(H0_planck)
    v = float(_R(H0)*_R(distance_Mpc))
    return {"velocity_km_s": v, "distance_Mpc": float(distance_Mpc),
            "H0": float(H0), "fraction_of_c": v/float(c_km_s)}


def calculate_redshift(velocity_km_s, relativistic=True):
    """
    Calculate cosmological redshift from recession velocity.

    Parameters
    ----------
    velocity_km_s : float  recession velocity in km/s
    relativistic  : bool   use relativistic formula if True

    Returns
    -------
    dict: redshift_z, velocity_km_s, wavelength_factor
    """
    beta = _R(velocity_km_s)/c_km_s
    if relativistic and abs(beta) < _R(1):
        z = float(np.sqrt((_R(1)+beta)/(_R(1)-beta)) - _R(1))
    else:
        z = float(beta)  # non-relativistic approximation
    return {"redshift_z": z, "velocity_km_s": float(velocity_km_s),
            "wavelength_factor": 1+z, "relativistic": relativistic}


def hubble_time(H0=None):
    """
    Hubble time t_H = 1/H0  (approximate age of universe).
    Returns time in Gyr.
    """
    if H0 is None: H0 = float(H0_planck)
    H0_per_s = _R(H0)*_R(1000)/Mpc_m   # s^-1
    t_s      = _R(1)/H0_per_s
    t_Gyr    = float(t_s/(_R(3.156e16)))  # Gyr
    return {"hubble_time_Gyr": t_Gyr, "H0": float(H0)}


def luminosity_distance(z, H0=None, Omega_m=0.315, Omega_L=0.685):
    """
    Luminosity distance d_L = (1+z) * d_C  using numerical integration.
    Returns distance in Mpc.
    """
    if H0 is None: H0 = float(H0_planck)
    # Comoving distance via Simpson quadrature
    n    = 500
    zs   = np.linspace(_R(0), _R(z), n, dtype=_R)
    dz   = zs[1]-zs[0]
    Ez   = np.sqrt(_R(Omega_m)*(1+zs)**3 + _R(Omega_L)).astype(_R)
    dC   = float(c_km_s/H0 * np.trapz(1/Ez, zs))
    dL   = dC*(1+z)
    return {"luminosity_distance_Mpc": dL, "comoving_distance_Mpc": dC, "z": z}


def distance_modulus(d_Mpc):
    """mu = 5*log10(d/10pc) = 5*log10(d_Mpc) + 25"""
    return float(5*np.log10(max(float(d_Mpc),1e-9)) + 25)


def scale_factor_evolution(t_arr, H0=None, Omega_m=0.315, Omega_L=0.685):
    """
    Approximate scale factor a(t) evolution (normalised a=1 at t=now).
    Simplified: uses a ~ (t/t0)^(2/3) for matter, (t/t0)^1 for Lambda.
    Returns a_arr float32.
    """
    if H0 is None: H0 = float(H0_planck)
    t_arr = np.asarray(t_arr, dtype=_R)
    t_H   = _R(hubble_time(H0)["hubble_time_Gyr"])
    a     = (_R(Omega_m)*(_R(1)-t_arr/t_H))**(_R(2/3))             + _R(Omega_L)*t_arr/t_H
    return np.clip(a, _R(0), None)
