
"""PKTron Cosmology — Dark Matter Density Profiles."""
import numpy as np
_R = np.float32

def nfw_profile(r, rho0, rs):
    """
    Navarro-Frenk-White dark matter halo profile.
    rho(r) = rho0 / [(r/rs)(1+r/rs)^2]
    """
    r  = np.asarray(r, dtype=_R)
    rs = _R(rs)
    return _R(rho0)/((r/rs)*(1+r/rs)**2)

def isothermal_profile(r, rho0, rc):
    """Singular isothermal sphere: rho = rho0/(1+(r/rc)^2)"""
    r = np.asarray(r, dtype=_R)
    return _R(rho0)/(1+(_R(r)/_R(rc))**2)

def dark_matter_fraction(r_arr, rho_fn, rho_total_fn):
    """
    Compute dark matter fraction f_DM(r) = rho_DM / rho_total at each r.
    """
    r     = np.asarray(r_arr, dtype=_R)
    rho_d = rho_fn(r)
    rho_t = rho_total_fn(r)
    return rho_d/np.maximum(rho_t, _R(1e-40))

def rotation_curve(r_arr, M_enclosed_fn, G=4.302e-3):
    """
    v_circ(r) = sqrt(G*M(r)/r)   G in pc*Msun^-1*(km/s)^2
    Returns velocity in km/s.
    """
    r = np.asarray(r_arr, dtype=_R)
    M = np.asarray([M_enclosed_fn(ri) for ri in r], dtype=_R)
    return np.sqrt(_R(G)*M/r)

def dark_matter_density_universe(z, Omega_dm=0.265, rho_c0=9.47e-27):
    """
    Cosmic DM density at redshift z: rho_DM(z) = Omega_DM * rho_c * (1+z)^3
    rho_c0 in kg/m^3  (critical density today).
    """
    return float(_R(Omega_dm)*_R(rho_c0)*(1+_R(z))**3)

def virial_mass(sigma_v, r_vir, G=6.674e-11):
    """M_vir = sigma^2 * r_vir / G  (virial theorem)"""
    return float(_R(sigma_v)**2*_R(r_vir)/_R(G))
