
"""PKTron Cosmology — Observational Cosmology helpers."""
import numpy as np
_R = np.float32

def angular_diameter_distance(z, H0=67.4, Omega_m=0.315, Omega_L=0.685):
    """d_A = d_C / (1+z)  in Mpc."""
    c_km_s = _R(299792.458)
    n    = 300
    zs   = np.linspace(_R(0), _R(z), n, dtype=_R)
    Ez   = np.sqrt(_R(Omega_m)*(1+zs)**3+_R(Omega_L))
    dC   = float(c_km_s/H0*np.trapz(1/Ez, zs))
    return {"angular_diameter_distance_Mpc": dC/(1+z),
            "comoving_distance_Mpc": dC, "z": z}

def cmb_temperature(z, T0=2.7255):
    """CMB temperature at redshift z: T(z) = T0*(1+z)"""
    return float(_R(T0)*(1+_R(z)))

def lookback_time(z, H0=67.4, Omega_m=0.315, Omega_L=0.685):
    """
    Lookback time in Gyr for given redshift z.
    """
    c_km_s = _R(299792.458)
    n      = 500
    zs     = np.linspace(_R(0), _R(z), n, dtype=_R)
    Ez     = np.sqrt(_R(Omega_m)*(1+zs)**3+_R(Omega_L))
    integrand = _R(1)/((1+zs)*Ez)
    t_H    = float(1/_R(H0)*_R(977.8))  # 1/H0 in Gyr
    t_lb   = float(t_H * np.trapz(integrand, zs))
    return {"lookback_time_Gyr": t_lb, "z": z}

def stellar_magnitude(flux, flux_ref=3.631e-23):
    """Apparent magnitude m = -2.5*log10(flux/flux_ref)  (AB system)."""
    return float(-2.5*np.log10(max(float(flux)/float(flux_ref), 1e-40)))

def absolute_magnitude(apparent_mag, distance_Mpc):
    """M = m - 5*log10(d_Mpc) - 25"""
    return float(apparent_mag - 5*np.log10(max(float(distance_Mpc),1e-9)) - 25)
