lcdm-core 0.1.0

// lcdm-core/src/params.rs

use crate::constants::{C, G, H0_UNITS_TO_INV_SEC, SIGMA_SB};
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

use crate::spec::{AccuracyPreset, DerivedParams, NeutrinoSpec};

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct ScientificParams {
    /// Canonical, user-facing parameter set (inputs in convenient forms).
    pub canonical: CanonicalParams,

    /// Explicit neutrino model specification (effective vs split).
    pub spec: NeutrinoSpec,

    /// Derived quantities computed from the canonical/spec inputs.
    pub derived: DerivedParams,

    /// Global accuracy / performance preset controlling numerical knobs.
    pub accuracy: AccuracyPreset,
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CanonicalParams {
    /// Dimensionless Hubble parameter h, where H0 = 100 h (km/s)/Mpc.
    pub h: f64,

    /// CMB temperature today in Kelvin.
    pub T_cmb0: f64,

    /// Physical baryon density: ω_b ≡ Ω_b h².
    pub omega_b: f64,

    /// Physical cold dark matter density: ω_cdm ≡ Ω_cdm h².
    pub omega_cdm: f64,

    /// Physical dark energy density: ω_Λ ≡ Ω_Λ h².
    pub omega_lambda: f64,

    /// Physical curvature density: ω_k ≡ Ω_k h².
    pub omega_k: f64,

    /// Effective number of relativistic species (or equivalent accounting in split mode).
    pub n_eff: f64,

    /// Neutrino masses in eV (interpretation depends on the neutrino spec).
    pub m_nu: Vec<f64>,

    /// Helium mass fraction (placeholder / future BBN integration).
    pub Y_He: f64,

    /// Dark energy equation-of-state parameter w0 (CPL form).
    pub w0: f64,

    /// Dark energy equation-of-state running parameter wa (CPL form).
    pub wa: f64,
}

impl CanonicalParams {
    /// Hubble constant H0 in (km/s)/Mpc.
    pub fn H0(&self) -> f64 {
        self.h * 100.0
    }
}

#[derive(Debug, Clone, Default)]
pub struct CosmoBuilder {
    h: Option<f64>,
    h0: Option<f64>,

    // Physical densities (ω = Ω h^2).
    omega_b_phys: Option<f64>,
    omega_cdm_phys: Option<f64>,

    // Dimensionless densities (Ω).
    Omega_b: Option<f64>,
    Omega_cdm: Option<f64>,
    Omega_m: Option<f64>,
    Omega_k: Option<f64>,
    Omega_L: Option<f64>,

    T_cmb0: Option<f64>,

    // Explicit neutrino model specification.
    nu_model: Option<NeutrinoSpec>,

    w0: Option<f64>,
    wa: Option<f64>,
    accuracy: AccuracyPreset,
    flat: bool,
}

impl CosmoBuilder {
    /// Create a new builder with default accuracy settings.
    pub fn new() -> Self {
        Self {
            accuracy: AccuracyPreset::Default,
            ..Self::default()
        }
    }

    /// Enforce spatial flatness (Ω_k = 0) and solve for ω_Λ by closure.
    pub fn flat(mut self) -> Self {
        self.flat = true;
        self
    }

    /// Set the dimensionless Hubble parameter h (H0 = 100 h).
    pub fn h(mut self, val: f64) -> Self {
        self.h = Some(val);
        self
    }

    /// Set H0 in (km/s)/Mpc (converted internally to h).
    pub fn H0(mut self, val: f64) -> Self {
        self.h0 = Some(val);
        self
    }

    // -------------------------------------------------------------------------
    // Dimensionless density API (Ω)
    // -------------------------------------------------------------------------

    #[allow(non_snake_case)]
    /// Set total matter density Ω_m.
    pub fn Omega_m(mut self, val: f64) -> Self {
        self.Omega_m = Some(val);
        self
    }

    #[allow(non_snake_case)]
    /// Set baryon density Ω_b.
    pub fn Omega_b(mut self, val: f64) -> Self {
        self.Omega_b = Some(val);
        self
    }

    #[allow(non_snake_case)]
    /// Set CDM density Ω_cdm.
    pub fn Omega_cdm(mut self, val: f64) -> Self {
        self.Omega_cdm = Some(val);
        self
    }

    #[allow(non_snake_case)]
    /// Set curvature density Ω_k.
    pub fn Omega_k(mut self, val: f64) -> Self {
        self.Omega_k = Some(val);
        self
    }

    #[allow(non_snake_case)]
    /// Set dark energy density Ω_Λ.
    pub fn Omega_L(mut self, val: f64) -> Self {
        self.Omega_L = Some(val);
        self
    }

    // -------------------------------------------------------------------------
    // Physical density API (ω ≡ Ω h²)
    // -------------------------------------------------------------------------

    /// Set physical baryon density ω_b.
    pub fn omega_b_phys(mut self, val: f64) -> Self {
        self.omega_b_phys = Some(val);
        self
    }

    /// Set physical CDM density ω_cdm.
    pub fn omega_cdm_phys(mut self, val: f64) -> Self {
        self.omega_cdm_phys = Some(val);
        self
    }

    // -------------------------------------------------------------------------
    // Misc configuration
    // -------------------------------------------------------------------------

    #[allow(non_snake_case)]
    /// Set the present-day CMB temperature T_cmb0 in Kelvin.
    pub fn T_cmb0(mut self, val: f64) -> Self {
        self.T_cmb0 = Some(val);
        self
    }

    /// Set the global accuracy preset controlling numerical settings.
    pub fn accuracy(mut self, val: AccuracyPreset) -> Self {
        self.accuracy = val;
        self
    }

    // -------------------------------------------------------------------------
    // Neutrino model selection
    // -------------------------------------------------------------------------

    /// Use an "effective" neutrino model with total `n_eff` and a mass list.
    ///
    /// For massless neutrinos, pass `[0.0]` (or multiple zeros if you want explicit states).
    pub fn neutrino_effective(mut self, n_eff: f64, masses_ev: Vec<f64>) -> Self {
        self.nu_model = Some(NeutrinoSpec::Effective { n_eff, masses_ev });
        self
    }

    /// Use a "split" neutrino model with explicit ultra-relativistic contribution `n_ur`
    /// plus a list of massive eigenstates and per-species temperature factors.
    pub fn neutrino_split(mut self, n_ur: f64, masses_ev: Vec<f64>, temp_factors: Vec<f64>) -> Self {
        self.nu_model = Some(NeutrinoSpec::Split {
            n_ur,
            masses_ev,
            temp_factors,
        });
        self
    }

    /// Build `ScientificParams` by resolving inputs, enforcing consistency, and computing
    /// all derived quantities.
    pub fn build(self) -> Result<ScientificParams, String> {
        // Resolve h from either h or H0 input (and validate if both were provided).
        let h = match (self.h, self.h0) {
            (Some(hv), None) => hv,
            (None, Some(h0v)) => h0v / 100.0,
            (Some(hv), Some(h0v)) => {
                if (hv * 100.0 - h0v).abs() > 1e-5 {
                    return Err("Inconsistent H0/h".into());
                }
                hv
            }
            (None, None) => return Err("Specify h or H0".into()),
        };

        // Default to the commonly used CMB temperature if not specified.
        let t_cmb0 = self.T_cmb0.unwrap_or(2.7255);

        // Neutrino model is mandatory (explicit).
        let spec = self.nu_model.ok_or("Neutrino model must be explicitly set")?;

        // Photon energy density today: ρ_γ = (4σ/c) T^4 (from blackbody radiation).
        let rho_g_energy = 4.0 * SIGMA_SB / C * t_cmb0.powi(4);

        // Critical density for h=1, using H100 = 100 (km/s)/Mpc converted to SI.
        let h100_si = 100.0 * H0_UNITS_TO_INV_SEC;
        let rho_crit_100 = 3.0 * h100_si * h100_si / (8.0 * std::f64::consts::PI * G);

        // Photon physical density ω_γ0 = ρ_γ / ρ_crit(h=1).
        let omega_g = (rho_g_energy / (C * C)) / rho_crit_100;

        // Base neutrino temperature: T_ν0 = T_γ0 (4/11)^(1/3).
        let t_nu_base = t_cmb0 * (4.0 / 11.0_f64).powf(1.0 / 3.0);

        // Numerical knobs for neutrino integration.
        let y_max = self.accuracy.neutrino_y_max();
        let steps = self.accuracy.neutrino_integ_steps();

        // Resolve ω_ur0 (explicit UR component), ω_ν0 (integrated massive/effective),
        // and the effective neutrino temperature used in the calculation.
        let (omega_ur0, omega_nu0, t_nu_eff) = match &spec {
            NeutrinoSpec::Effective { n_eff, masses_ev } => {
                // For effective mode, masses must be provided (use [0.0] for massless).
                if masses_ev.is_empty() {
                    return Err(
                        "Effective neutrino model: masses_ev must be non-empty (use [0.0] for massless)"
                            .into(),
                    );
                }

                // Distribute N_eff across the provided number of states via an effective temperature.
                let n = masses_ev.len().max(1) as f64;
                let t_eff = t_nu_base * (n_eff / n).powf(0.25);

                // Hardcode degeneracy_g = 2.0 for each species.
                let o_nu = crate::neutrino::calc_omega_nu_z_cfg(
                    masses_ev,
                    t_eff,
                    0.0,
                    y_max,
                    steps,
                    2.0,
                );
                (0.0, o_nu, t_eff)
            }
            NeutrinoSpec::Split {
                n_ur,
                masses_ev,
                temp_factors,
            } => {
                // Validate split-model inputs.
                if *n_ur < 0.0 {
                    return Err("n_ur < 0".into());
                }
                if masses_ev.len() != temp_factors.len() {
                    return Err("Length mismatch".into());
                }
                if masses_ev.iter().any(|&m| m < 1e-10) {
                    return Err("Split massive must be > 0".into());
                }
                if temp_factors.iter().any(|&tf| tf <= 0.0) {
                    return Err("Temp factors must be positive".into());
                }

                // Ultra-relativistic contribution uses the standard N_eff scaling:
                //   ρ_ν,UR / ρ_γ = (7/8) * (4/11)^(4/3) * N_eff
                let o_ur =
                    *n_ur * (7.0 / 8.0) * (4.0 / 11.0_f64).powf(4.0 / 3.0) * omega_g;

                // Integrate each massive eigenstate with its temperature factor.
                let mut o_nu_total = 0.0;
                for (&m, &tf) in masses_ev.iter().zip(temp_factors.iter()) {
                    o_nu_total += crate::neutrino::calc_omega_nu_z_cfg(
                        &[m],
                        t_nu_base * tf,
                        0.0,
                        y_max,
                        steps,
                        2.0,
                    );
                }

                (o_ur, o_nu_total, t_nu_base)
            }
        };

        // ---------------------------------------------------------------------
        // Resolve baryon density ω_b
        // ---------------------------------------------------------------------

        let omega_b = if let (Some(wb), Some(Ob)) = (self.omega_b_phys, self.Omega_b) {
            // If both are provided, enforce consistency.
            let wb2 = Ob * h * h;
            if (wb - wb2).abs() > 1e-10 {
                return Err("Inconsistent omega_b_phys vs Omega_b".into());
            }
            wb
        } else if let Some(wb) = self.omega_b_phys {
            wb
        } else if let Some(Ob) = self.Omega_b {
            Ob * h * h
        } else {
            return Err("Specify Omega_b or omega_b_phys".into());
        };

        // ---------------------------------------------------------------------
        // Resolve CDM density ω_cdm
        // ---------------------------------------------------------------------

        let omega_cdm = if let (Some(wc), Some(Oc)) = (self.omega_cdm_phys, self.Omega_cdm) {
            // If both are provided, enforce consistency.
            let wc2 = Oc * h * h;
            if (wc - wc2).abs() > 1e-10 {
                return Err("Inconsistent omega_cdm_phys vs Omega_cdm".into());
            }
            wc
        } else if let Some(wc) = self.omega_cdm_phys {
            wc
        } else if let Some(Oc) = self.Omega_cdm {
            Oc * h * h
        } else if let Some(Om) = self.Omega_m {
            // If only Ω_m is given, infer ω_cdm = Ω_m h^2 - ω_b.
            Om * h * h - omega_b
        } else {
            return Err("Specify Omega_cdm, omega_cdm_phys, or Omega_m".into());
        };

        // Prevent nonsensical negative CDM (e.g., Ω_m < Ω_b).
        if omega_cdm < 0.0 {
            return Err("Negative CDM (Omega_m < Omega_b)".into());
        }

        let omega_m = omega_b + omega_cdm;
        let omega_non_de = omega_g + omega_ur0 + omega_nu0 + omega_m;

        // ---------------------------------------------------------------------
        // Resolve dark energy and curvature via closure
        // ---------------------------------------------------------------------

        let (omega_lambda, omega_k) = if self.flat {
            // In flat mode, Ω_k must be 0 and ω_Λ is solved by closure:
            //   h^2 = ω_non_de + ω_Λ
            if let Some(ok) = self.Omega_k {
                if ok.abs() > 1e-12 {
                    return Err("flat() is set but Omega_k was provided (must be 0)".into());
                }
            }

            let ol = h * h - omega_non_de;
            if ol < 0.0 {
                return Err("Total density > h^2 in flat model".into());
            }

            (ol, 0.0)
        } else {
            // Non-flat: allow explicit Ω_Λ and/or Ω_k if provided, otherwise solve by closure.
            let ol = if let Some(o_l_in) = self.Omega_L {
                let ol_sq = o_l_in * h * h;

                // If Ω_k is also provided, ensure the residual is consistent.
                if let Some(o_k_in) = self.Omega_k {
                    let ok_sq = o_k_in * h * h;
                    let residual = (h * h - omega_non_de - ol_sq - ok_sq).abs();
                    if residual > 1e-6 {
                        return Err("Inconsistent Omega_L and Omega_k".into());
                    }
                }

                ol_sq
            } else {
                // If Ω_Λ is not provided, infer it after subtracting any provided Ω_k.
                h * h - omega_non_de - self.Omega_k.unwrap_or(0.0) * h * h
            };

            // ω_k is what remains after enforcing closure: h^2 = ω_non_de + ω_Λ + ω_k.
            (ol, h * h - omega_non_de - ol)
        };

        // Build the derived-parameter cache.
        let derived = DerivedParams::new(
            h,
            t_cmb0,
            omega_b,
            omega_cdm,
            omega_k,
            omega_lambda,
            omega_g,
            omega_ur0,
            omega_nu0,
            t_nu_base,
            t_nu_eff,
        );

        // Canonical n_eff and m_nu bookkeeping:
        // - Effective: pass through n_eff and the provided mass list.
        // - Split: treat n_eff as (n_ur + number_of_massive_states).
        let (n_eff_can, m_nu_can) = match &spec {
            NeutrinoSpec::Effective { n_eff, masses_ev, .. } => (*n_eff, masses_ev.clone()),
            NeutrinoSpec::Split { n_ur, masses_ev, .. } => (*n_ur + masses_ev.len() as f64, masses_ev.clone()),
        };

        Ok(ScientificParams {
            canonical: CanonicalParams {
                h,
                T_cmb0: t_cmb0,
                omega_b,
                omega_cdm,
                omega_lambda,
                omega_k,
                n_eff: n_eff_can,
                m_nu: m_nu_can,
                Y_He: 0.24,
                w0: self.w0.unwrap_or(-1.0),
                wa: self.wa.unwrap_or(0.0),
            },
            spec,
            derived,
            accuracy: self.accuracy,
        })
    }
}