nufast 0.5.0

//! # NuFast
//!
//! Fast and accurate three-flavor neutrino oscillation probabilities in vacuum and matter.
//!
//! This is a Rust port of [NuFast](https://github.com/PeterDenton/NuFast) by Peter Denton,
//! based on the algorithm described in [arXiv:2405.02400](https://arxiv.org/abs/2405.02400).
//!
//! ## Features
//!
//! - **Vacuum oscillations**: Full 3-flavor PMNS oscillation probabilities
//! - **Matter effects (MSW)**: Constant-density matter with arbitrary electron fraction
//! - **Anti-neutrino mode**: Sign-flipped δCP and matter potential
//! - **CP violation**: Full δ_CP phase support
//! - **Newton refinement**: Optional iterative improvement for matter eigenvalues
//! - **Batch API**: `VacuumBatch` for 45% faster energy spectrum calculations
//! - **Zero dependencies**: Pure Rust, no external crates required
//! - **WebAssembly**: Zig implementation compiles to WASM (see `benchmarks/zig/wasm/`)
//!
//! ## Quick Start
//!
//! ```rust
//! use nufast::{VacuumParameters, probability_vacuum_lbl};
//! use std::f64::consts::PI;
//!
//! let params = VacuumParameters {
//!     s12sq: 0.307,
//!     s13sq: 0.0218,
//!     s23sq: 0.545,
//!     delta: 1.36 * PI,
//!     Dmsq21: 7.42e-5,
//!     Dmsq31: 2.517e-3,
//!     L: 295.0,
//!     E: 0.6,
//!     antineutrino: false,
//! };
//!
//! let probs = probability_vacuum_lbl(&params);
//! assert!(probs[1][0] >= 0.0 && probs[1][0] <= 1.0); // P(νμ → νe)
//! ```
//!
//! ## Performance
//!
//! Benchmarks on AMD Ryzen (10M iterations):
//!
//! | Language | Vacuum | Matter N=0 |
//! |----------|--------|------------|
//! | Zig SIMD | 21 ns  | 37 ns      |
//! | Zig      | 42 ns  | 108 ns     |
//! | **Rust** | 61 ns  | **95 ns**  |
//! | C++      | 49 ns  | 130 ns     |
//! | WASM     | 80 ns  | 150 ns     |
//!
//! Rust is **27% faster than C++** for matter calculations.
//! Zig SIMD achieves **48M vacuum/sec** throughput.

#![cfg_attr(feature = "no_std", no_std)]
// Allow physics naming conventions (standard in neutrino oscillation literature)
#![allow(non_snake_case)]

use core::f64::consts::PI;

/// Conversion factor: eV² × km → GeV (divided by 4)
/// 
/// This combines ħc and unit conversions for the oscillation phase:
/// Δ = Δm² × L / (4E) in natural units
const EV_SQ_KM_TO_GEV_OVER4: f64 = 1e-9 / 1.97327e-7 * 1e3 / 4.0;

/// Matter potential conversion factor: Y_e × ρ × E → A
/// 
/// A = √2 G_F N_e E where N_e = Y_e × ρ × N_A / m_nucleon
/// This constant is approximately 1.52 × 10⁻⁴ eV² / (g/cm³ × GeV)
const YE_RHO_E_TO_A: f64 = 1.52e-4;

/// Parameters for vacuum oscillation calculations.
///
/// All angles are specified as sin²θ, CP phase in radians.
/// Mass splittings in eV², baseline in km, energy in GeV.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct VacuumParameters {
    /// sin²θ₁₂ (solar mixing angle)
    pub s12sq: f64,
    /// sin²θ₁₃ (reactor mixing angle)
    pub s13sq: f64,
    /// sin²θ₂₃ (atmospheric mixing angle)
    pub s23sq: f64,
    /// CP-violating phase δ in radians
    pub delta: f64,
    /// Δm²₂₁ in eV² (solar mass splitting, always positive)
    pub Dmsq21: f64,
    /// Δm²₃₁ in eV² (atmospheric mass splitting, positive for NO, negative for IO)
    pub Dmsq31: f64,
    /// Baseline L in km
    pub L: f64,
    /// Neutrino energy E in GeV
    pub E: f64,
    /// Anti-neutrino mode (flips sign of δCP)
    pub antineutrino: bool,
}

impl VacuumParameters {
    /// Create parameters with NuFit 5.2 best-fit values (Normal Ordering).
    ///
    /// Only L and E need to be specified.
    pub fn nufit52_no(l: f64, e: f64) -> Self {
        Self {
            s12sq: 0.307,
            s13sq: 0.02203,
            s23sq: 0.546,
            delta: 1.36 * PI,
            Dmsq21: 7.42e-5,
            Dmsq31: 2.517e-3,
            L: l,
            E: e,
            antineutrino: false,
        }
    }

    /// Create parameters with NuFit 5.2 best-fit values (Inverted Ordering).
    pub fn nufit52_io(l: f64, e: f64) -> Self {
        Self {
            s12sq: 0.307,
            s13sq: 0.02219,
            s23sq: 0.539,
            delta: 1.56 * PI,
            Dmsq21: 7.42e-5,
            Dmsq31: -2.498e-3,
            L: l,
            E: e,
            antineutrino: false,
        }
    }
    
    /// Get effective delta (sign-flipped for antineutrinos)
    pub fn effective_delta(&self) -> f64 {
        if self.antineutrino { -self.delta } else { self.delta }
    }
}

/// Parameters for matter oscillation calculations.
///
/// Extends [VacuumParameters] with matter density and electron fraction.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MatterParameters {
    /// sin²θ₁₂ (solar mixing angle)
    pub s12sq: f64,
    /// sin²θ₁₃ (reactor mixing angle)
    pub s13sq: f64,
    /// sin²θ₂₃ (atmospheric mixing angle)
    pub s23sq: f64,
    /// CP-violating phase δ in radians
    pub delta: f64,
    /// Δm²₂₁ in eV² (solar mass splitting)
    pub Dmsq21: f64,
    /// Δm²₃₁ in eV² (atmospheric mass splitting)
    pub Dmsq31: f64,
    /// Baseline L in km
    pub L: f64,
    /// Neutrino energy E in GeV
    pub E: f64,
    /// Matter density ρ in g/cm³
    pub rho: f64,
    /// Electron fraction Y_e (typically 0.5 for Earth)
    pub Ye: f64,
    /// Number of Newton iterations for eigenvalue refinement (0 = DMP approximation)
    pub N_Newton: u8,
    /// Anti-neutrino mode (flips sign of matter potential and δCP)
    pub antineutrino: bool,
}

impl MatterParameters {
    /// Create parameters with NuFit 5.2 best-fit values (Normal Ordering).
    ///
    /// Uses Earth crust density (2.6 g/cm³) and Y_e = 0.5.
    pub fn nufit52_no(l: f64, e: f64) -> Self {
        Self {
            s12sq: 0.307,
            s13sq: 0.02203,
            s23sq: 0.546,
            delta: 1.36 * PI,
            Dmsq21: 7.42e-5,
            Dmsq31: 2.517e-3,
            L: l,
            E: e,
            rho: 2.6,
            Ye: 0.5,
            N_Newton: 0,
            antineutrino: false,
        }
    }

    /// Create from vacuum parameters, adding matter properties.
    pub fn from_vacuum(vac: &VacuumParameters, rho: f64, ye: f64, n_newton: u8) -> Self {
        Self {
            s12sq: vac.s12sq,
            s13sq: vac.s13sq,
            s23sq: vac.s23sq,
            delta: vac.delta,
            Dmsq21: vac.Dmsq21,
            Dmsq31: vac.Dmsq31,
            L: vac.L,
            E: vac.E,
            rho,
            Ye: ye,
            N_Newton: n_newton,
            antineutrino: vac.antineutrino,
        }
    }
    
    /// Get effective delta (sign-flipped for antineutrinos)
    pub fn effective_delta(&self) -> f64 {
        if self.antineutrino { -self.delta } else { self.delta }
    }
    
    /// Get matter potential sign (negative for antineutrinos)
    pub fn matter_sign(&self) -> f64 {
        if self.antineutrino { -1.0 } else { 1.0 }
    }
}

/// 3×3 probability matrix type.
///
/// Indexed as `probs[α][β]` = P(ν_α → ν_β) where:
/// - 0 = electron (e)
/// - 1 = muon (μ)
/// - 2 = tau (τ)
pub type ProbabilityMatrix = [[f64; 3]; 3];

/// Calculate vacuum oscillation probabilities for long-baseline experiments.
///
/// Returns a 3×3 matrix where `probs[α][β]` = P(ν_α → ν_β).
///
/// # Example
///
/// ```rust
/// use nufast::{VacuumParameters, probability_vacuum_lbl};
///
/// let params = VacuumParameters::nufit52_no(295.0, 0.6);
/// let probs = probability_vacuum_lbl(&params);
///
/// // Check unitarity
/// let row_sum: f64 = probs[1].iter().sum();
/// assert!((row_sum - 1.0).abs() < 1e-10);
/// ```
pub fn probability_vacuum_lbl(parameters: &VacuumParameters) -> ProbabilityMatrix {
    let VacuumParameters {
        s12sq,
        s13sq,
        s23sq,
        delta: _,
        Dmsq21,
        Dmsq31,
        L,
        E,
        antineutrino: _,
    } = *parameters;
    
    // Use effective delta (sign-flipped for antineutrinos)
    let effective_delta = parameters.effective_delta();

    // Pre-calculate trig functions
    let c13sq = 1.0 - s13sq;

    // Ueisq's
    let Ue3sq = s13sq;
    let Ue2sq = c13sq * s12sq;

    // Umisq's, Utisq's and Jvac
    let Um3sq = c13sq * s23sq;
    // Um2sq and Ut2sq are used here as temporary variables
    let Ut2sq = s13sq * s12sq * s23sq;
    let Um2sq = (1.0 - s12sq) * (1.0 - s23sq);

    let Jrr = (Um2sq * Ut2sq).sqrt();
    let sind = effective_delta.sin();
    let cosd = effective_delta.cos();
    let Um2sq = Um2sq + Ut2sq - 2.0 * Jrr * cosd;
    let Jvac = 8.0 * Jrr * c13sq * sind;

    // Get all elements of Usq using unitarity
    let Ue1sq = 1.0 - Ue3sq - Ue2sq;
    let Um1sq = 1.0 - Um3sq - Um2sq;

    let Ut3sq = 1.0 - Um3sq - Ue3sq;
    let Ut2sq = 1.0 - Um2sq - Ue2sq;
    let Ut1sq = 1.0 - Um1sq - Ue1sq;

    // Get the kinematic terms
    let Lover4E = EV_SQ_KM_TO_GEV_OVER4 * L / E;

    let D21 = Dmsq21 * Lover4E;
    let D31 = Dmsq31 * Lover4E;

    let sinD21 = D21.sin();
    let sinD31 = D31.sin();
    let sinD32 = (D31 - D21).sin();

    let triple_sin = sinD21 * sinD31 * sinD32;

    let sinsqD21_2 = 2.0 * sinD21 * sinD21;
    let sinsqD31_2 = 2.0 * sinD31 * sinD31;
    let sinsqD32_2 = 2.0 * sinD32 * sinD32;

    // Calculate the three probabilities, separating CPC and CPV
    let Pme_CPC = (Ut3sq - Um2sq * Ue1sq - Um1sq * Ue2sq) * sinsqD21_2
        + (Ut2sq - Um3sq * Ue1sq - Um1sq * Ue3sq) * sinsqD31_2
        + (Ut1sq - Um3sq * Ue2sq - Um2sq * Ue3sq) * sinsqD32_2;

    let Pme_CPV = -Jvac * triple_sin;

    let Pmm = 1.0
        - 2.0
            * (Um2sq * Um1sq * sinsqD21_2
                + Um3sq * Um1sq * sinsqD31_2
                + Um3sq * Um2sq * sinsqD32_2);

    let Pee = 1.0
        - 2.0
            * (Ue2sq * Ue1sq * sinsqD21_2
                + Ue3sq * Ue1sq * sinsqD31_2
                + Ue3sq * Ue2sq * sinsqD32_2);

    // Assign all the probabilities
    let mut probs = [[0.0; 3]; 3];

    probs[0][0] = Pee; // Pee
    probs[0][1] = Pme_CPC - Pme_CPV; // Pem
    probs[0][2] = 1.0 - Pee - probs[0][1]; // Pet

    probs[1][0] = Pme_CPC + Pme_CPV; // Pme
    probs[1][1] = Pmm; // Pmm
    probs[1][2] = 1.0 - probs[1][0] - Pmm; // Pmt

    probs[2][0] = 1.0 - Pee - probs[1][0]; // Pte
    probs[2][1] = 1.0 - probs[0][1] - Pmm; // Ptm
    probs[2][2] = 1.0 - probs[0][2] - probs[1][2]; // Ptt

    probs
}

/// Calculate matter oscillation probabilities for long-baseline experiments.
///
/// Uses the DMP approximation with optional Newton refinement for the
/// matter-modified eigenvalues.
///
/// # Example
///
/// ```rust
/// use nufast::{MatterParameters, probability_matter_lbl};
///
/// let params = MatterParameters::nufit52_no(1300.0, 2.5);
/// let probs = probability_matter_lbl(&params);
///
/// // Matter enhances νμ → νe for normal ordering
/// assert!(probs[1][0] > 0.0);
/// ```
pub fn probability_matter_lbl(parameters: &MatterParameters) -> ProbabilityMatrix {
    let MatterParameters {
        s12sq,
        s13sq,
        s23sq,
        delta: _,
        Dmsq21,
        Dmsq31,
        L,
        E,
        rho,
        Ye,
        N_Newton,
        antineutrino: _,
    } = *parameters;
    
    // Use effective delta and matter sign (both flipped for antineutrinos)
    let effective_delta = parameters.effective_delta();
    let matter_sign = parameters.matter_sign();

    // Pre-calculate trig functions
    let c13sq = 1.0 - s13sq;

    // Ueisq's
    let Ue2sq = c13sq * s12sq;
    let Ue3sq = s13sq;

    // Umisq's, Utisq's and Jmatter
    let Um3sq = c13sq * s23sq;
    let Ut2sq = s13sq * s12sq * s23sq;
    let Um2sq = (1.0 - s12sq) * (1.0 - s23sq);

    let Jrr = (Um2sq * Ut2sq).sqrt();
    let sind = effective_delta.sin();
    let cosd = effective_delta.cos();

    let Um2sq = Um2sq + Ut2sq - 2.0 * Jrr * cosd;
    let Jmatter = 8.0 * Jrr * c13sq * sind;
    let Amatter = matter_sign * Ye * rho * E * YE_RHO_E_TO_A;
    let Dmsqee = Dmsq31 - s12sq * Dmsq21;

    // Calculate A, B, C, See, Tee, and part of Tmm
    let A_sum = Dmsq21 + Dmsq31; // temporary variable
    let See = A_sum - Dmsq21 * Ue2sq - Dmsq31 * Ue3sq;
    let Tmm_base = Dmsq21 * Dmsq31;
    let Tee = Tmm_base * (1.0 - Ue3sq - Ue2sq);
    let C = Amatter * Tee;
    let A = A_sum + Amatter;

    // Get lambda3 from lambda+ of MP/DMP
    let xmat = Amatter / Dmsqee;
    let tmp = 1.0 - xmat;
    let mut lambda3 = Dmsq31 + 0.5 * Dmsqee * (xmat - 1.0 + (tmp * tmp + 4.0 * s13sq * xmat).sqrt());

    // Newton iterations to improve lambda3 arbitrarily
    let B = Tmm_base + Amatter * See; // B is only needed for N_Newton >= 1
    for _ in 0..N_Newton {
        lambda3 = (lambda3 * lambda3 * (lambda3 - A) + C)
            / (lambda3 * (2.0 * lambda3 - A) + B);
    }

    // Get Delta lambdas
    let tmp = A - lambda3;
    let Dlambda21 = (tmp * tmp - 4.0 * C / lambda3).sqrt();
    let lambda2 = 0.5 * (A - lambda3 + Dlambda21);
    let Dlambda32 = lambda3 - lambda2;
    let Dlambda31 = Dlambda32 + Dlambda21;

    // Use Rosetta for Veisq's (matter-modified PMNS elements)
    let PiDlambdaInv = 1.0 / (Dlambda31 * Dlambda32 * Dlambda21);
    let Xp3 = PiDlambdaInv * Dlambda21;
    let Xp2 = -PiDlambdaInv * Dlambda31;

    let Ue3sq = (lambda3 * (lambda3 - See) + Tee) * Xp3;
    let Ue2sq = (lambda2 * (lambda2 - See) + Tee) * Xp2;

    let Smm = A - Dmsq21 * Um2sq - Dmsq31 * Um3sq;
    let Tmm = Tmm_base * (1.0 - Um3sq - Um2sq) + Amatter * (See + Smm - A_sum);

    let Um3sq = (lambda3 * (lambda3 - Smm) + Tmm) * Xp3;
    let Um2sq = (lambda2 * (lambda2 - Smm) + Tmm) * Xp2;

    // Use NHS for J (matter-modified Jarlskog)
    let Jmatter = Jmatter * Dmsq21 * Dmsq31 * (Dmsq31 - Dmsq21) * PiDlambdaInv;

    // Get all elements of Usq
    let Ue1sq = 1.0 - Ue3sq - Ue2sq;
    let Um1sq = 1.0 - Um3sq - Um2sq;

    let Ut3sq = 1.0 - Um3sq - Ue3sq;
    let Ut2sq = 1.0 - Um2sq - Ue2sq;
    let Ut1sq = 1.0 - Um1sq - Ue1sq;

    // Get the kinematic terms (using matter eigenvalues)
    let Lover4E = EV_SQ_KM_TO_GEV_OVER4 * L / E;

    let D21 = Dlambda21 * Lover4E;
    let D32 = Dlambda32 * Lover4E;

    let sinD21 = D21.sin();
    let sinD31 = (D32 + D21).sin();
    let sinD32 = D32.sin();

    let triple_sin = sinD21 * sinD31 * sinD32;

    let sinsqD21_2 = 2.0 * sinD21 * sinD21;
    let sinsqD31_2 = 2.0 * sinD31 * sinD31;
    let sinsqD32_2 = 2.0 * sinD32 * sinD32;

    // Calculate the three necessary probabilities
    let Pme_CPC = (Ut3sq - Um2sq * Ue1sq - Um1sq * Ue2sq) * sinsqD21_2
        + (Ut2sq - Um3sq * Ue1sq - Um1sq * Ue3sq) * sinsqD31_2
        + (Ut1sq - Um3sq * Ue2sq - Um2sq * Ue3sq) * sinsqD32_2;
    let Pme_CPV = -Jmatter * triple_sin;

    let Pmm = 1.0
        - 2.0
            * (Um2sq * Um1sq * sinsqD21_2
                + Um3sq * Um1sq * sinsqD31_2
                + Um3sq * Um2sq * sinsqD32_2);

    let Pee = 1.0
        - 2.0
            * (Ue2sq * Ue1sq * sinsqD21_2
                + Ue3sq * Ue1sq * sinsqD31_2
                + Ue3sq * Ue2sq * sinsqD32_2);

    // Assign all the probabilities
    let mut probs = [[0.0; 3]; 3];

    probs[0][0] = Pee;
    probs[0][1] = Pme_CPC - Pme_CPV;
    probs[0][2] = 1.0 - Pee - probs[0][1];

    probs[1][0] = Pme_CPC + Pme_CPV;
    probs[1][1] = Pmm;
    probs[1][2] = 1.0 - probs[1][0] - Pmm;

    probs[2][0] = 1.0 - Pee - probs[1][0];
    probs[2][1] = 1.0 - probs[0][1] - Pmm;
    probs[2][2] = 1.0 - probs[0][2] - probs[1][2];

    probs
}

/// Pre-computed vacuum oscillation parameters for batch calculations.
/// 
/// Use this when computing probabilities over many energy or baseline values.
/// The mixing matrix elements are computed once, reducing per-call overhead by ~30%.
/// 
/// # Example
/// 
/// ```rust
/// use nufast::VacuumBatch;
/// use std::f64::consts::PI;
/// 
/// let batch = VacuumBatch::new(0.307, 0.0218, 0.545, 1.36 * PI, 7.42e-5, 2.517e-3, false);
/// 
/// // Compute probabilities at many energies
/// for e in [0.5, 1.0, 2.0, 3.0, 5.0] {
///     let probs = batch.probability_at(295.0, e);
///     println!("P(νμ→νe) at E={} GeV: {:.4}", e, probs[1][0]);
/// }
/// ```
#[derive(Debug, Clone, Copy)]
pub struct VacuumBatch {
    // Pre-computed mixing matrix elements
    Ue1sq: f64,
    Ue2sq: f64,
    Ue3sq: f64,
    Um1sq: f64,
    Um2sq: f64,
    Um3sq: f64,
    Ut1sq: f64,
    Ut2sq: f64,
    Ut3sq: f64,
    
    // CP violation term
    Jvac: f64,
    
    // Mass splittings
    Dmsq21: f64,
    Dmsq31: f64,
}

impl VacuumBatch {
    /// Create a batch calculator with pre-computed mixing elements.
    /// 
    /// For antineutrino mode, pass the negative of delta (or use `from_params`).
    pub fn new(
        s12sq: f64,
        s13sq: f64,
        s23sq: f64,
        delta: f64,
        Dmsq21: f64,
        Dmsq31: f64,
        antineutrino: bool,
    ) -> Self {
        let c13sq = 1.0 - s13sq;
        let effective_delta = if antineutrino { -delta } else { delta };

        let Ue3sq = s13sq;
        let Ue2sq = c13sq * s12sq;

        let Um3sq = c13sq * s23sq;
        let Ut2sq_temp = s13sq * s12sq * s23sq;
        let Um2sq_temp = (1.0 - s12sq) * (1.0 - s23sq);

        let Jrr = (Um2sq_temp * Ut2sq_temp).sqrt();
        let sind = effective_delta.sin();
        let cosd = effective_delta.cos();
        let Um2sq = Um2sq_temp + Ut2sq_temp - 2.0 * Jrr * cosd;
        let Jvac = 8.0 * Jrr * c13sq * sind;

        let Ue1sq = 1.0 - Ue3sq - Ue2sq;
        let Um1sq = 1.0 - Um3sq - Um2sq;

        let Ut3sq = 1.0 - Um3sq - Ue3sq;
        let Ut2sq = 1.0 - Um2sq - Ue2sq;
        let Ut1sq = 1.0 - Um1sq - Ue1sq;

        Self {
            Ue1sq, Ue2sq, Ue3sq,
            Um1sq, Um2sq, Um3sq,
            Ut1sq, Ut2sq, Ut3sq,
            Jvac,
            Dmsq21, Dmsq31,
        }
    }
    
    /// Create from VacuumParameters (handles antineutrino mode automatically).
    pub fn from_params(params: &VacuumParameters) -> Self {
        Self::new(
            params.s12sq,
            params.s13sq,
            params.s23sq,
            params.delta,
            params.Dmsq21,
            params.Dmsq31,
            params.antineutrino,
        )
    }
    
    /// Create from NuFit 5.2 best-fit values (Normal Ordering).
    pub fn nufit52_no() -> Self {
        Self::new(0.307, 0.02203, 0.546, 1.36 * PI, 7.42e-5, 2.517e-3, false)
    }
    
    /// Create from NuFit 5.2 best-fit values (Inverted Ordering).
    pub fn nufit52_io() -> Self {
        Self::new(0.307, 0.02219, 0.539, 1.56 * PI, 7.42e-5, -2.498e-3, false)
    }
    
    /// Compute probability matrix at given baseline and energy.
    /// 
    /// This is ~30% faster than `probability_vacuum_lbl` when called repeatedly.
    #[inline]
    pub fn probability_at(&self, l: f64, e: f64) -> ProbabilityMatrix {
        let Lover4E = EV_SQ_KM_TO_GEV_OVER4 * l / e;

        let D21 = self.Dmsq21 * Lover4E;
        let D31 = self.Dmsq31 * Lover4E;

        let sinD21 = D21.sin();
        let sinD31 = D31.sin();
        let sinD32 = (D31 - D21).sin();

        let triple_sin = sinD21 * sinD31 * sinD32;

        let sinsqD21_2 = 2.0 * sinD21 * sinD21;
        let sinsqD31_2 = 2.0 * sinD31 * sinD31;
        let sinsqD32_2 = 2.0 * sinD32 * sinD32;

        let Pme_CPC = (self.Ut3sq - self.Um2sq * self.Ue1sq - self.Um1sq * self.Ue2sq) * sinsqD21_2
            + (self.Ut2sq - self.Um3sq * self.Ue1sq - self.Um1sq * self.Ue3sq) * sinsqD31_2
            + (self.Ut1sq - self.Um3sq * self.Ue2sq - self.Um2sq * self.Ue3sq) * sinsqD32_2;

        let Pme_CPV = -self.Jvac * triple_sin;

        let Pmm = 1.0
            - 2.0
                * (self.Um2sq * self.Um1sq * sinsqD21_2
                    + self.Um3sq * self.Um1sq * sinsqD31_2
                    + self.Um3sq * self.Um2sq * sinsqD32_2);

        let Pee = 1.0
            - 2.0
                * (self.Ue2sq * self.Ue1sq * sinsqD21_2
                    + self.Ue3sq * self.Ue1sq * sinsqD31_2
                    + self.Ue3sq * self.Ue2sq * sinsqD32_2);

        let mut probs = [[0.0; 3]; 3];

        probs[0][0] = Pee;
        probs[0][1] = Pme_CPC - Pme_CPV;
        probs[0][2] = 1.0 - Pee - probs[0][1];

        probs[1][0] = Pme_CPC + Pme_CPV;
        probs[1][1] = Pmm;
        probs[1][2] = 1.0 - probs[1][0] - Pmm;

        probs[2][0] = 1.0 - Pee - probs[1][0];
        probs[2][1] = 1.0 - probs[0][1] - Pmm;
        probs[2][2] = 1.0 - probs[0][2] - probs[1][2];

        probs
    }
    
    /// Compute probabilities for a single initial flavor over many energies.
    /// 
    /// This is the most efficient way to compute energy spectra.
    /// Returns a Vec of (energy, Pe, Pmu, Ptau) tuples.
    #[cfg(not(feature = "no_std"))]
    pub fn spectrum(&self, l: f64, energies: &[f64], initial_flavor: usize) -> Vec<(f64, f64, f64, f64)> {
        energies.iter().map(|&e| {
            let probs = self.probability_at(l, e);
            let row = initial_flavor.min(2);
            (e, probs[row][0], probs[row][1], probs[row][2])
        }).collect()
    }
    
    /// Compute probabilities along a baseline at fixed energy.
    #[cfg(not(feature = "no_std"))]
    pub fn baseline_scan(&self, e: f64, baselines: &[f64], initial_flavor: usize) -> Vec<(f64, f64, f64, f64)> {
        baselines.iter().map(|&l| {
            let probs = self.probability_at(l, e);
            let row = initial_flavor.min(2);
            (l, probs[row][0], probs[row][1], probs[row][2])
        }).collect()
    }
}

/// Clamp all probabilities to [0, 1] and ensure row unitarity.
/// 
/// Useful when numerical precision issues cause slight violations.
pub fn normalize_probabilities(probs: &mut ProbabilityMatrix) {
    for row in probs.iter_mut() {
        for p in row.iter_mut() {
            *p = p.clamp(0.0, 1.0);
        }
        let sum: f64 = row.iter().sum();
        if sum > 0.0 && (sum - 1.0).abs() > 1e-10 {
            for p in row.iter_mut() {
                *p /= sum;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const EPSILON: f64 = 1e-10;

    #[test]
    fn test_vacuum_unitarity() {
        let params = VacuumParameters::nufit52_no(1000.0, 2.0);
        let probs = probability_vacuum_lbl(&params);

        // Check all rows sum to 1
        for (i, row) in probs.iter().enumerate() {
            let sum: f64 = row.iter().sum();
            assert!(
                (sum - 1.0).abs() < EPSILON,
                "Row {} sum = {}, expected 1.0",
                i,
                sum
            );
        }
    }

    #[test]
    fn test_vacuum_zero_distance() {
        let params = VacuumParameters::nufit52_no(0.0, 2.0);
        let probs = probability_vacuum_lbl(&params);

        // At L=0, no oscillation: diagonal = 1, off-diagonal = 0
        for i in 0..3 {
            for j in 0..3 {
                let expected = if i == j { 1.0 } else { 0.0 };
                assert!(
                    (probs[i][j] - expected).abs() < EPSILON,
                    "probs[{}][{}] = {}, expected {}",
                    i,
                    j,
                    probs[i][j],
                    expected
                );
            }
        }
    }

    #[test]
    fn test_matter_unitarity() {
        let params = MatterParameters::nufit52_no(1300.0, 2.5);
        let mut probs = probability_matter_lbl(&params);
        normalize_probabilities(&mut probs);

        for (i, row) in probs.iter().enumerate() {
            let sum: f64 = row.iter().sum();
            assert!(
                (sum - 1.0).abs() < 1e-6,
                "Row {} sum = {}, expected 1.0",
                i,
                sum
            );
        }
    }

    #[test]
    fn test_matter_enhances_appearance() {
        // For normal ordering, matter enhances νμ → νe at DUNE energies
        let vac = VacuumParameters::nufit52_no(1300.0, 2.5);
        let mat = MatterParameters::nufit52_no(1300.0, 2.5);

        let probs_vac = probability_vacuum_lbl(&vac);
        let probs_mat = probability_matter_lbl(&mat);

        // Matter should modify the appearance probability
        assert!(
            (probs_mat[1][0] - probs_vac[1][0]).abs() > 0.001,
            "Matter effect should modify P(νμ → νe)"
        );
    }

    #[test]
    fn test_cp_conjugate() {
        // P(ν, δ) should differ from P(ν, -δ) due to CP violation
        let mut params_pos = VacuumParameters::nufit52_no(1000.0, 2.0);
        params_pos.delta = 1.36 * PI;

        let mut params_neg = params_pos;
        params_neg.delta = -1.36 * PI;

        let probs_pos = probability_vacuum_lbl(&params_pos);
        let probs_neg = probability_vacuum_lbl(&params_neg);

        // P(μ→e, δ) ≠ P(μ→e, -δ) unless δ = 0 or π
        assert!(
            (probs_pos[1][0] - probs_neg[1][0]).abs() > 1e-6,
            "CP conjugate probabilities should differ"
        );
    }
    
    #[test]
    fn test_batch_vacuum_consistency() {
        let batch = VacuumBatch::new(0.307, 0.0218, 0.545, 1.36 * PI, 7.42e-5, 2.517e-3, false);
        let energies = [0.5, 1.0, 2.0, 3.0, 5.0];
        let l = 295.0;
        
        for &e in &energies {
            let probs_batch = batch.probability_at(l, e);
            let probs_single = probability_vacuum_lbl(&VacuumParameters {
                s12sq: 0.307,
                s13sq: 0.0218,
                s23sq: 0.545,
                delta: 1.36 * PI,
                Dmsq21: 7.42e-5,
                Dmsq31: 2.517e-3,
                L: l,
                E: e,
                antineutrino: false,
            });
            
            for i in 0..3 {
                for j in 0..3 {
                    assert!(
                        (probs_batch[i][j] - probs_single[i][j]).abs() < EPSILON,
                        "Batch and single differ at E={}: [{},{}] {} vs {}",
                        e, i, j, probs_batch[i][j], probs_single[i][j]
                    );
                }
            }
        }
    }
    
    #[test]
    fn test_antineutrino_vacuum() {
        let l = 1300.0;
        let e = 2.5;
        
        let nu_params = VacuumParameters::nufit52_no(l, e);
        let mut nubar_params = nu_params;
        nubar_params.antineutrino = true;
        
        let nu_probs = probability_vacuum_lbl(&nu_params);
        let nubar_probs = probability_vacuum_lbl(&nubar_params);
        
        // Pee and Pmm should be the same (no CP violation in diagonal channels)
        assert!(
            (nu_probs[0][0] - nubar_probs[0][0]).abs() < EPSILON,
            "Pee should be the same for ν and ν̄"
        );
        assert!(
            (nu_probs[1][1] - nubar_probs[1][1]).abs() < EPSILON,
            "Pmm should be the same for ν and ν̄"
        );
        
        // Pme(ν) should equal Pem(ν̄) due to CPT theorem
        assert!(
            (nu_probs[1][0] - nubar_probs[0][1]).abs() < EPSILON,
            "CPT: P(νμ→νe) should equal P(ν̄e→ν̄μ)"
        );
        assert!(
            (nu_probs[0][1] - nubar_probs[1][0]).abs() < EPSILON,
            "CPT: P(νe→νμ) should equal P(ν̄μ→ν̄e)"
        );
    }
    
    #[test]
    fn test_antineutrino_matter() {
        let l = 1300.0;
        let e = 2.5;
        
        let nu_params = MatterParameters::nufit52_no(l, e);
        let mut nubar_params = nu_params;
        nubar_params.antineutrino = true;
        
        let nu_probs = probability_matter_lbl(&nu_params);
        let nubar_probs = probability_matter_lbl(&nubar_params);
        
        // In matter, probabilities should differ due to opposite matter potential
        let diff = (nu_probs[1][0] - nubar_probs[1][0]).abs();
        assert!(
            diff > 0.01,
            "Matter effect should differ for ν and ν̄: diff = {}",
            diff
        );
    }
}