//! # nereids-python
//!
//! PyO3 Python bindings for the NEREIDS neutron resonance imaging library.
//!
//! Provides a Pythonic API for:
//! - Computing theoretical transmission spectra
//! - Fitting measured transmission to recover isotopic compositions
//! - Spatial mapping across imaging data
//!
//! ## Usage
//! ```python
//! import nereids
//! import numpy as np
//!
//! # Load ENDF data for U-238
//! isotope = nereids.load_endf(92, 238)
//!
//! # Compute transmission spectrum
//! energies = np.linspace(1.0, 30.0, 1000)
//! transmission = nereids.forward_model(energies, [(isotope, 0.001)], temperature_k=300.0)
//!
//! # Fit a measured spectrum
//! result = nereids.fit_spectrum(measured_t, sigma, energies, [isotope])
//! ```

use std::sync::Arc;

use numpy::{
    PyArray1, PyArray2, PyArray3, PyReadonlyArray1, PyReadonlyArray2, PyReadonlyArray3,
    PyUntypedArrayMethods,
};
use pyo3::prelude::*;

use nereids_core::elements;
use nereids_core::types::Isotope;
use nereids_endf::parser::parse_endf_file2;
use nereids_endf::resonance::{
    LGroup, Resonance, ResonanceData, ResonanceFormalism, ResonanceRange,
};
use nereids_endf::retrieval::{EndfLibrary, EndfRetriever, mat_number};
use nereids_fitting::lm::{self, FitModel, LmConfig};
use nereids_fitting::parameters::{FitParameter, ParameterSet};
use nereids_fitting::poisson::{self, CountsModel, PoissonConfig, TemperatureContext};
use nereids_fitting::transmission_model::{PrecomputedTransmissionModel, TransmissionFitModel};
use nereids_io::normalization::{self as norm, NormalizationParams};
use nereids_io::tof::BeamlineParams;
use nereids_physics::doppler::{self, DopplerParams};
use nereids_physics::resolution::{
    self, ResolutionFunction, ResolutionParams, TabulatedResolution,
};
use nereids_physics::transmission::{self, InstrumentParams, SampleParams};
use nereids_pipeline::detectability;
use nereids_pipeline::pipeline::FitConfig;
use nereids_pipeline::sparse::{SparseConfig, estimate_nuisance, sparse_reconstruct};

/// Python wrapper for ENDF resonance data.
///
/// Uses `Arc` internally so that `.clone()` in `py.detach()` closures is O(1)
/// (refcount bump) instead of deep-copying the entire resonance dataset.
#[pyclass(name = "ResonanceData", from_py_object)]
#[derive(Clone)]
struct PyResonanceData {
    inner: Arc<ResonanceData>,
}

#[pymethods]
impl PyResonanceData {
    /// String representation.
    fn __repr__(&self) -> String {
        let n_res: usize = self
            .inner
            .ranges
            .iter()
            .flat_map(|r| &r.l_groups)
            .map(|lg| lg.resonances.len())
            .sum();
        format!(
            "ResonanceData(Z={}, A={}, AWR={:.3}, n_resonances={})",
            self.inner.isotope.z(),
            self.inner.isotope.a(),
            self.inner.awr,
            n_res
        )
    }

    /// Atomic number.
    #[getter]
    fn z(&self) -> u32 {
        self.inner.isotope.z()
    }

    /// Mass number.
    #[getter]
    fn a(&self) -> u32 {
        self.inner.isotope.a()
    }

    /// Atomic weight ratio.
    #[getter]
    fn awr(&self) -> f64 {
        self.inner.awr
    }

    /// Number of resonances.
    #[getter]
    fn n_resonances(&self) -> usize {
        self.inner
            .ranges
            .iter()
            .flat_map(|r| &r.l_groups)
            .map(|lg| lg.resonances.len())
            .sum()
    }

    /// Target spin (I) of the first resonance range.
    #[getter]
    fn target_spin(&self) -> f64 {
        self.inner
            .ranges
            .first()
            .map(|r| r.target_spin)
            .unwrap_or(0.0)
    }

    /// Effective scattering radius (fm).
    ///
    /// Returns the global AP from the first range. If AP=0 (common in
    /// ENDF Reich-Moore data that uses energy-dependent radii), falls back
    /// to the first L-group's channel radius APL.
    #[getter]
    fn scattering_radius(&self) -> f64 {
        self.inner
            .ranges
            .first()
            .map(|r| {
                if r.scattering_radius != 0.0 {
                    r.scattering_radius
                } else {
                    // Fall back to first L-group's channel radius
                    r.l_groups.first().map(|lg| lg.apl).unwrap_or(0.0)
                }
            })
            .unwrap_or(0.0)
    }

    /// Orbital angular momentum values (L) present in the data.
    #[getter]
    fn l_values(&self) -> Vec<u32> {
        let mut ls: Vec<u32> = self
            .inner
            .ranges
            .iter()
            .flat_map(|r| &r.l_groups)
            .map(|lg| lg.l)
            .collect();
        ls.sort();
        ls.dedup();
        ls
    }
}

/// Result of fitting a spectrum.
#[pyclass(name = "FitResult")]
struct PyFitResult {
    densities: Vec<f64>,
    /// `None` when covariance computation was skipped.
    uncertainties: Option<Vec<f64>>,
    reduced_chi_squared: f64,
    converged: bool,
    iterations: usize,
    /// Fitted temperature in Kelvin (only meaningful when `fit_temperature=True`).
    temperature_k: Option<f64>,
    /// 1-sigma uncertainty on fitted temperature (K).
    temperature_k_unc: Option<f64>,
}

#[pymethods]
impl PyFitResult {
    /// Fitted areal densities (atoms/barn).
    #[getter]
    fn densities<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_vec(py, self.densities.clone())
    }

    /// Uncertainties on fitted densities.
    ///
    /// Returns NaN-filled array when covariance computation was skipped.
    #[getter]
    fn uncertainties<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        let unc = self
            .uncertainties
            .clone()
            .unwrap_or_else(|| vec![f64::NAN; self.densities.len()]);
        PyArray1::from_vec(py, unc)
    }

    /// Reduced chi-squared of the fit.
    #[getter]
    fn reduced_chi_squared(&self) -> f64 {
        self.reduced_chi_squared
    }

    /// Whether the fit converged.
    #[getter]
    fn converged(&self) -> bool {
        self.converged
    }

    /// Number of iterations.
    #[getter]
    fn iterations(&self) -> usize {
        self.iterations
    }

    /// Fitted sample temperature in Kelvin (``None`` when ``fit_temperature=False``).
    #[getter]
    fn temperature_k(&self) -> Option<f64> {
        self.temperature_k
    }

    /// 1-sigma uncertainty on fitted temperature in Kelvin (``None`` when
    /// ``fit_temperature=False``).
    #[getter]
    fn temperature_k_unc(&self) -> Option<f64> {
        self.temperature_k_unc
    }

    fn __repr__(&self) -> String {
        if let Some(t) = self.temperature_k {
            format!(
                "FitResult(converged={}, chi2_red={:.4}, densities={:?}, temperature_k={:.1})",
                self.converged, self.reduced_chi_squared, self.densities, t
            )
        } else {
            format!(
                "FitResult(converged={}, chi2_red={:.4}, densities={:?})",
                self.converged, self.reduced_chi_squared, self.densities
            )
        }
    }
}

/// Python wrapper for tabulated resolution function.
///
/// Uses `Arc` internally so that `.clone()` in `py.detach()` closures is O(1).
#[pyclass(name = "TabulatedResolution", from_py_object)]
#[derive(Clone)]
struct PyTabulatedResolution {
    inner: Arc<TabulatedResolution>,
}

#[pymethods]
impl PyTabulatedResolution {
    /// Number of reference energies.
    #[getter]
    fn n_energies(&self) -> usize {
        self.inner.ref_energies().len()
    }

    /// Energy range (min, max) of the reference kernels in eV.
    #[getter]
    fn energy_range(&self) -> (f64, f64) {
        let e = self.inner.ref_energies();
        if e.is_empty() {
            (0.0, 0.0)
        } else {
            (e[0], e[e.len() - 1])
        }
    }

    /// Flight path length in meters.
    #[getter]
    fn flight_path_m(&self) -> f64 {
        self.inner.flight_path_m()
    }

    /// Number of points per kernel.
    #[getter]
    fn points_per_kernel(&self) -> usize {
        self.inner
            .kernels()
            .first()
            .map(|(o, _)| o.len())
            .unwrap_or(0)
    }

    fn __repr__(&self) -> String {
        let (lo, hi) = self.energy_range();
        format!(
            "TabulatedResolution(n_energies={}, range=[{:.4e}, {:.4e}] eV, flight_path={:.1} m)",
            self.n_energies(),
            lo,
            hi,
            self.inner.flight_path_m(),
        )
    }
}

/// Result of spatial (per-pixel) mapping.
///
/// Numpy arrays are constructed once and cached; property access returns
/// cheap references (refcount bump) rather than copying data.
#[pyclass(name = "SpatialResult")]
struct PySpatialResult {
    density_maps: Vec<Py<PyArray2<f64>>>,
    uncertainty_maps: Vec<Py<PyArray2<f64>>>,
    chi_squared_map: Py<PyArray2<f64>>,
    converged_map: Py<PyArray2<bool>>,
    n_converged: usize,
    n_total: usize,
    isotope_names: Vec<String>,
    shape: (usize, usize),
}

#[pymethods]
impl PySpatialResult {
    /// Density maps as a list of 2D numpy arrays, one per isotope.
    #[getter]
    fn density_maps<'py>(&self, py: Python<'py>) -> Vec<Bound<'py, PyArray2<f64>>> {
        self.density_maps
            .iter()
            .map(|m| m.bind(py).clone())
            .collect()
    }

    /// Uncertainty maps as a list of 2D numpy arrays.
    #[getter]
    fn uncertainty_maps<'py>(&self, py: Python<'py>) -> Vec<Bound<'py, PyArray2<f64>>> {
        self.uncertainty_maps
            .iter()
            .map(|m| m.bind(py).clone())
            .collect()
    }

    /// Reduced chi-squared map.
    #[getter]
    fn chi_squared_map<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<f64>> {
        self.chi_squared_map.bind(py).clone()
    }

    /// Convergence map (True = converged).
    #[getter]
    fn converged_map<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<bool>> {
        self.converged_map.bind(py).clone()
    }

    /// Number of converged pixels.
    #[getter]
    fn n_converged(&self) -> usize {
        self.n_converged
    }

    /// Total number of fitted pixels.
    #[getter]
    fn n_total(&self) -> usize {
        self.n_total
    }

    /// Isotope names.
    #[getter]
    fn isotope_names(&self) -> Vec<String> {
        self.isotope_names.clone()
    }

    fn __repr__(&self) -> String {
        format!(
            "SpatialResult(shape={}x{}, isotopes={}, converged={}/{})",
            self.shape.0,
            self.shape.1,
            self.isotope_names.len(),
            self.n_converged,
            self.n_total,
        )
    }
}

/// Python wrapper for sparse (Poisson/KL) reconstruction results.
///
/// Returned by `spatial_map(..., fitter='poisson')`.
/// Uses raw count data and open-beam for statistically optimal fitting
/// at all count levels including very low statistics (< ~10 counts/bin).
#[pyclass(name = "SparseResult")]
struct PySparseResult {
    density_maps: Vec<Py<PyArray2<f64>>>,
    nll_map: Py<PyArray2<f64>>,
    converged_map: Py<PyArray2<bool>>,
    n_converged: usize,
    n_total: usize,
    isotope_names: Vec<String>,
    shape: (usize, usize),
}

#[pymethods]
impl PySparseResult {
    /// Density maps as a list of 2D numpy arrays, one per isotope.
    #[getter]
    fn density_maps<'py>(&self, py: Python<'py>) -> Vec<Bound<'py, PyArray2<f64>>> {
        self.density_maps
            .iter()
            .map(|m| m.bind(py).clone())
            .collect()
    }

    /// Poisson negative log-likelihood map (analogous to chi_squared_map for LM).
    #[getter]
    fn nll_map<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<f64>> {
        self.nll_map.bind(py).clone()
    }

    /// Convergence map (True = converged).
    #[getter]
    fn converged_map<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<bool>> {
        self.converged_map.bind(py).clone()
    }

    /// Number of converged pixels.
    #[getter]
    fn n_converged(&self) -> usize {
        self.n_converged
    }

    /// Total number of fitted pixels.
    #[getter]
    fn n_total(&self) -> usize {
        self.n_total
    }

    /// Isotope names.
    #[getter]
    fn isotope_names(&self) -> Vec<String> {
        self.isotope_names.clone()
    }

    fn __repr__(&self) -> String {
        format!(
            "SparseResult(shape={}x{}, isotopes={}, converged={}/{})",
            self.shape.0,
            self.shape.1,
            self.isotope_names.len(),
            self.n_converged,
            self.n_total,
        )
    }
}

/// Compute cross-sections at given energies for an isotope.
///
/// Args:
///     energies: Energy grid in eV (1D numpy array).
///     data: ResonanceData for the isotope.
///
/// Returns:
///     Dictionary with 'total', 'elastic', 'capture', 'fission' arrays.
#[pyfunction]
fn cross_sections<'py>(
    py: Python<'py>,
    energies: PyReadonlyArray1<f64>,
    data: &PyResonanceData,
) -> PyResult<Bound<'py, pyo3::types::PyDict>> {
    let e_owned = energies.as_slice()?.to_vec();
    let res_data = data.inner.clone();

    // Release the GIL for the cross-section computation.
    // Use cross_sections_on_grid() which precomputes J-group data once,
    // rather than recomputing per energy point via cross_sections_at_energy().
    let (total, elastic, capture, fission) = py.detach(move || {
        let results = nereids_physics::reich_moore::cross_sections_on_grid(&res_data, &e_owned);
        let mut total = Vec::with_capacity(results.len());
        let mut elastic = Vec::with_capacity(results.len());
        let mut capture = Vec::with_capacity(results.len());
        let mut fission = Vec::with_capacity(results.len());
        for xs in results {
            total.push(xs.total);
            elastic.push(xs.elastic);
            capture.push(xs.capture);
            fission.push(xs.fission);
        }
        (total, elastic, capture, fission)
    });

    let dict = pyo3::types::PyDict::new(py);
    dict.set_item("total", PyArray1::from_vec(py, total))?;
    dict.set_item("elastic", PyArray1::from_vec(py, elastic))?;
    dict.set_item("capture", PyArray1::from_vec(py, capture))?;
    dict.set_item("fission", PyArray1::from_vec(py, fission))?;
    Ok(dict)
}

/// Compute theoretical transmission spectrum.
///
/// Resolution broadening can be applied via either Gaussian parameters
/// (``flight_path_m``, ``delta_t_us``, ``delta_l_m``) or a tabulated
/// resolution function (``resolution``). Providing both is an error.
///
/// Args:
///     energies: Energy grid in eV (1D numpy array).
///     isotopes: List of (ResonanceData, areal_density) tuples.
///     temperature_k: Sample temperature in Kelvin (default 0.0).
///     flight_path_m: Flight path in meters for Gaussian resolution (optional).
///     delta_t_us: Timing uncertainty in microseconds (optional).
///     delta_l_m: Path length uncertainty in meters (optional).
///     delta_e_us: Exponential tail parameter in SAMMY Deltae units (optional,
///         default None/0.0). When non-zero, adds an exponential tail to the
///         resolution kernel (SAMMY Iesopr=3).
///     resolution: TabulatedResolution from ``load_resolution()`` (optional).
///
/// Returns:
///     1D numpy array of transmission values.
#[pyfunction]
#[pyo3(signature = (energies, isotopes, temperature_k=0.0, flight_path_m=None, delta_t_us=None, delta_l_m=None, resolution=None, delta_e_us=None))]
fn forward_model<'py>(
    py: Python<'py>,
    energies: PyReadonlyArray1<f64>,
    isotopes: Vec<(PyResonanceData, f64)>,
    temperature_k: f64,
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let e_owned = energies.as_slice()?.to_vec();

    let sample_isotopes: Vec<(ResonanceData, f64)> = isotopes
        .into_iter()
        .map(|(d, thick)| (Arc::unwrap_or_clone(d.inner), thick))
        .collect();

    let sample = SampleParams::new(temperature_k, sample_isotopes)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

    let res_fn = build_resolution(flight_path_m, delta_t_us, delta_l_m, resolution, delta_e_us)?;
    let instrument = res_fn.map(|r| InstrumentParams { resolution: r });

    // Release the GIL for the forward model computation.
    let t = py.detach(move || transmission::forward_model(&e_owned, &sample, instrument.as_ref()));
    let t = t.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
    Ok(PyArray1::from_vec(py, t))
}

/// Fit a measured transmission spectrum to recover isotopic areal densities.
///
/// Resolution broadening can be applied via either Gaussian parameters
/// (``flight_path_m``, ``delta_t_us``, ``delta_l_m``) or a tabulated
/// resolution function (``resolution``). Providing both is an error.
///
/// Args:
///     measured_t: Measured transmission (1D numpy array).
///     sigma: Measurement uncertainties (1D numpy array).
///     energies: Energy grid in eV (1D numpy array).
///     isotopes: List of ResonanceData objects.
///     temperature_k: Sample temperature in Kelvin (default 0.0).
///     initial_densities: Initial guess for areal densities (optional).
///     max_iter: Maximum LM iterations (default 100).
///     flight_path_m: Flight path in meters for Gaussian resolution (optional).
///     delta_t_us: Timing uncertainty in microseconds (optional).
///     delta_l_m: Path length uncertainty in meters (optional).
///     resolution: TabulatedResolution from ``load_resolution()`` (optional).
///     fit_temperature: If ``True``, treat ``temperature_k`` as the initial
///         guess and include it as a free parameter in the fit.
///         The fitted value and 1-sigma uncertainty are returned in
///         ``result.temperature_k`` and ``result.temperature_k_unc``.
///     fitter: Optimizer to use.
///
///         - ``"lm"`` (default): Levenberg-Marquardt (chi-squared).
///           ``measured_t`` is transmission, ``sigma`` is uncertainties.
///         - ``"poisson"``: Poisson NLL with analytical gradient.
///           ``measured_t`` is sample counts, ``sigma`` is open-beam/flux counts.
///           ``result.reduced_chi_squared`` contains the final NLL (not chi-squared).
///
/// Returns:
///     FitResult with densities, uncertainties, and fit quality.
#[pyfunction]
#[pyo3(signature = (measured_t, sigma, energies, isotopes, temperature_k=0.0, initial_densities=None, max_iter=100, flight_path_m=None, delta_t_us=None, delta_l_m=None, resolution=None, delta_e_us=None, fit_temperature=false, fitter="lm"))]
fn fit_spectrum(
    py: Python<'_>,
    measured_t: PyReadonlyArray1<f64>,
    sigma: PyReadonlyArray1<f64>,
    energies: PyReadonlyArray1<f64>,
    isotopes: Vec<PyResonanceData>,
    temperature_k: f64,
    initial_densities: Option<Vec<f64>>,
    max_iter: usize,
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
    fit_temperature: bool,
    fitter: &str,
) -> PyResult<PyFitResult> {
    let e = energies.as_slice()?;
    let t = measured_t.as_slice()?;
    let s = sigma.as_slice()?;

    if t.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "measured_t must not be empty",
        ));
    }
    if t.len() != s.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "measured_t length ({}) must match sigma length ({})",
            t.len(),
            s.len(),
        )));
    }
    if t.len() != e.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "measured_t length ({}) must match energies length ({})",
            t.len(),
            e.len(),
        )));
    }
    if isotopes.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "isotopes list must not be empty",
        ));
    }

    let n_isotopes = isotopes.len();
    let res_data: Vec<ResonanceData> = isotopes
        .into_iter()
        .map(|d| Arc::unwrap_or_clone(d.inner))
        .collect();

    let init = initial_densities.unwrap_or_else(|| vec![0.001; n_isotopes]);

    if init.len() != n_isotopes {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "initial_densities length ({}) must match isotopes length ({})",
            init.len(),
            n_isotopes,
        )));
    }

    if !temperature_k.is_finite() || temperature_k < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temperature_k must be finite and non-negative",
        ));
    }
    if fit_temperature && temperature_k < 1.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temperature_k must be >= 1.0 when fit_temperature=True \
             (it is used as the initial guess; the optimizer lower bound is 1 K)",
        ));
    }

    let res_fn = build_resolution(flight_path_m, delta_t_us, delta_l_m, resolution, delta_e_us)?;

    match fitter {
        "lm" => {
            let instrument = res_fn.map(|r| Arc::new(InstrumentParams { resolution: r }));

            // Copy numpy slices to owned vectors so we can release the GIL.
            let e_owned = e.to_vec();
            let t_owned = t.to_vec();
            let s_owned = s.to_vec();

            // Release the GIL for the heavy computation.
            // The closure uses only Rust types; PyErr conversion happens outside.
            let result: Result<PyFitResult, String> = py.detach(move || {
                let mut param_vec: Vec<FitParameter> = init
                    .iter()
                    .enumerate()
                    .map(|(i, &d)| FitParameter::non_negative(format!("isotope_{}", i), d))
                    .collect();

                if fit_temperature {
                    param_vec.push(FitParameter {
                        name: "temperature_k".into(),
                        value: temperature_k,
                        lower: 1.0,
                        upper: 5000.0,
                        fixed: false,
                    });
                }

                let mut params = ParameterSet::new(param_vec);
                let config = LmConfig {
                    max_iter,
                    ..LmConfig::default()
                };

                // When temperature is fixed, precompute Doppler+resolution-broadened
                // cross-sections once and use PrecomputedTransmissionModel for the LM
                // fit.  This avoids recomputing XS on every model evaluation and also
                // provides the analytical Jacobian for free.
                // When temperature is free, cross-sections change every iteration so
                // we must use the full TransmissionFitModel.
                let lm_result = if fit_temperature {
                    let model = TransmissionFitModel::new(
                        e_owned,
                        res_data,
                        temperature_k,
                        instrument,
                        (0..n_isotopes).collect(),
                        Some(n_isotopes),
                        None,
                    )
                    .map_err(|e| format!("TransmissionFitModel::new failed: {e}"))?;
                    lm::levenberg_marquardt(&model, &t_owned, &s_owned, &mut params, &config)
                        .map_err(|e| format!("levenberg_marquardt failed: {e}"))?
                } else {
                    let xs = transmission::broadened_cross_sections(
                        &e_owned,
                        &res_data,
                        temperature_k,
                        instrument.as_ref().map(Arc::as_ref),
                        None,
                    )
                    .map_err(|e| format!("broadened_cross_sections failed: {e}"))?;
                    let density_indices: Arc<Vec<usize>> = Arc::new((0..n_isotopes).collect());
                    let precomputed = PrecomputedTransmissionModel {
                        cross_sections: Arc::new(xs),
                        density_indices,
                    };
                    lm::levenberg_marquardt(&precomputed, &t_owned, &s_owned, &mut params, &config)
                        .map_err(|e| format!("levenberg_marquardt failed: {e}"))?
                };

                let densities: Vec<f64> = (0..n_isotopes).map(|i| lm_result.params[i]).collect();

                let (uncertainties, fitted_temperature, fitted_temperature_unc) = match lm_result
                    .uncertainties
                {
                    Some(unc) => {
                        let n_total = n_isotopes + if fit_temperature { 1 } else { 0 };
                        if unc.len() < n_total {
                            return Err(format!(
                                "uncertainty vector length ({}) should be >= parameter count ({})",
                                unc.len(),
                                n_total
                            ));
                        }
                        let density_unc: Vec<f64> = (0..n_isotopes)
                            .map(|i| unc.get(i).copied().unwrap_or(f64::NAN))
                            .collect();
                        let (t, t_unc) = if fit_temperature {
                            (
                                Some(lm_result.params[n_isotopes]),
                                Some(unc.get(n_isotopes).copied().unwrap_or(f64::NAN)),
                            )
                        } else {
                            (None, None)
                        };
                        (Some(density_unc), t, t_unc)
                    }
                    None => {
                        // Covariance was skipped — no uncertainties.
                        let t = if fit_temperature {
                            Some(lm_result.params[n_isotopes])
                        } else {
                            None
                        };
                        (None, t, None)
                    }
                };

                Ok(PyFitResult {
                    densities,
                    uncertainties,
                    reduced_chi_squared: lm_result.reduced_chi_squared,
                    converged: lm_result.converged,
                    iterations: lm_result.iterations,
                    temperature_k: fitted_temperature,
                    temperature_k_unc: fitted_temperature_unc,
                })
            });
            result.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e))
        }

        "poisson" => {
            // Copy numpy slices to owned vectors so we can release the GIL.
            let e_owned = e.to_vec();
            let y_obs = t.to_vec();
            let flux_owned = s.to_vec();

            let instrument = res_fn.map(|r| InstrumentParams { resolution: r });

            // Release the GIL for the heavy computation.
            // The closure uses only Rust types; PyErr conversion happens outside.
            let result: Result<PyFitResult, String> = py.detach(move || {
                // When fitting temperature, compute unbroadened XS first (Reich-Moore,
                // temperature-independent), then derive broadened XS from them.
                // This avoids computing Reich-Moore twice.
                let base_xs_arc = if fit_temperature {
                    Some(Arc::new(
                        transmission::unbroadened_cross_sections(&e_owned, &res_data, None)
                            .map_err(|e| format!("unbroadened_cross_sections failed: {e}"))?,
                    ))
                } else {
                    None
                };

                // Precompute broadened cross-sections at the initial temperature.
                let xs = if let Some(ref base) = base_xs_arc {
                    transmission::broadened_cross_sections_from_base(
                        &e_owned,
                        base,
                        &res_data,
                        temperature_k,
                        instrument.as_ref(),
                    )
                    .map_err(|e| format!("broadened_cross_sections_from_base failed: {e}"))?
                } else {
                    transmission::broadened_cross_sections(
                        &e_owned,
                        &res_data,
                        temperature_k,
                        instrument.as_ref(),
                        None,
                    )
                    .map_err(|e| format!("broadened_cross_sections failed: {e}"))?
                };

                // Build density parameters.
                let mut param_vec: Vec<FitParameter> = init
                    .iter()
                    .enumerate()
                    .map(|(i, &d)| FitParameter::non_negative(format!("isotope_{}", i), d))
                    .collect();

                // When fitting temperature, add it as an extra free parameter.
                if fit_temperature {
                    param_vec.push(FitParameter {
                        name: "temperature_k".into(),
                        value: temperature_k,
                        lower: 1.0,
                        upper: 5000.0,
                        fixed: false,
                    });
                }

                let mut params = ParameterSet::new(param_vec);

                // The transmission model for CountsModel: uses precomputed xs
                // (or full physics model when fitting temperature).
                let density_indices: Arc<Vec<usize>> = Arc::new((0..n_isotopes).collect());
                let xs_arc = Arc::new(xs.clone());
                let precomputed = PrecomputedTransmissionModel {
                    cross_sections: xs_arc,
                    density_indices: Arc::clone(&density_indices),
                };

                // When fitting temperature, share base_xs between
                // TemperatureContext and TransmissionFitModel.
                let temp_ctx = if fit_temperature {
                    Some(TemperatureContext {
                        temperature_index: n_isotopes,
                        resonance_data: res_data.clone(),
                        energies: e_owned.clone(),
                        instrument: instrument.clone(),
                        base_xs: base_xs_arc.clone(),
                    })
                } else {
                    None
                };

                let full_model;
                let t_model: &dyn FitModel = if fit_temperature {
                    full_model = TransmissionFitModel::new(
                        e_owned,
                        res_data,
                        temperature_k,
                        instrument.map(Arc::new),
                        (*density_indices).clone(),
                        Some(n_isotopes),
                        base_xs_arc,
                    )
                    .map_err(|e| format!("TransmissionFitModel::new failed: {e}"))?;
                    &full_model
                } else {
                    &precomputed
                };

                let background = vec![0.0f64; flux_owned.len()];
                let counts_model = CountsModel {
                    transmission_model: t_model,
                    flux: &flux_owned,
                    background: &background,
                };

                let poisson_config = PoissonConfig {
                    max_iter,
                    ..PoissonConfig::default()
                };

                let poisson_result = poisson::poisson_fit_analytic(
                    &counts_model,
                    &y_obs,
                    &flux_owned,
                    &xs,
                    &density_indices,
                    &mut params,
                    &poisson_config,
                    temp_ctx.as_ref(),
                )
                .map_err(|e| format!("poisson_fit_analytic failed: {e}"))?;

                let densities: Vec<f64> =
                    (0..n_isotopes).map(|i| poisson_result.params[i]).collect();
                // Poisson optimizer does not compute uncertainties from covariance;
                // report None (Python getter will return NaN-filled array).
                let uncertainties: Option<Vec<f64>> = None;

                let (fitted_temperature, fitted_temperature_unc) = if fit_temperature {
                    (Some(poisson_result.params[n_isotopes]), Some(f64::NAN))
                } else {
                    (None, None)
                };

                Ok(PyFitResult {
                    densities,
                    uncertainties,
                    reduced_chi_squared: poisson_result.nll,
                    converged: poisson_result.converged,
                    iterations: poisson_result.iterations,
                    temperature_k: fitted_temperature,
                    temperature_k_unc: fitted_temperature_unc,
                })
            });
            result.map_err(|e| pyo3::exceptions::PyValueError::new_err(e))
        }

        other => Err(pyo3::exceptions::PyValueError::new_err(format!(
            "fitter must be 'lm' or 'poisson', got '{}'",
            other,
        ))),
    }
}

/// Convert time-of-flight (μs) to energy (eV).
///
/// Args:
///     tof_us: Time-of-flight in microseconds.
///     flight_path_m: Flight path in meters.
///
/// Returns:
///     Energy in eV.
#[pyfunction]
fn tof_to_energy(tof_us: f64, flight_path_m: f64) -> f64 {
    nereids_core::constants::tof_to_energy(tof_us, flight_path_m)
}

/// Convert energy (eV) to time-of-flight (μs).
///
/// Args:
///     energy_ev: Energy in eV.
///     flight_path_m: Flight path in meters.
///
/// Returns:
///     Time-of-flight in microseconds.
#[pyfunction]
fn energy_to_tof(energy_ev: f64, flight_path_m: f64) -> f64 {
    nereids_core::constants::energy_to_tof(energy_ev, flight_path_m)
}

/// Load ENDF resonance data for an isotope from the IAEA database.
///
/// Downloads and parses the ENDF file, caching it locally at
/// ``~/.cache/nereids/endf/`` for subsequent calls.
///
/// Args:
///     z: Atomic number (e.g. 92 for uranium).
///     a: Mass number (e.g. 238).
///     library: ENDF library name. One of "endf8.0", "endf8.1" (default),
///              "jeff3.3", "jendl5".
///     mat: ENDF MAT (material) number. If None, looks up from built-in table
///          (~40 common isotopes). Provide explicitly for uncommon isotopes.
///
/// Returns:
///     ResonanceData parsed from the ENDF file.
#[pyfunction]
#[pyo3(signature = (z, a, library="endf8.1", mat=None))]
fn load_endf(
    py: Python<'_>,
    z: u32,
    a: u32,
    library: &str,
    mat: Option<u32>,
) -> PyResult<PyResonanceData> {
    let lib = match library {
        "endf8.0" | "endf/b-viii.0" => EndfLibrary::EndfB8_0,
        "endf8.1" | "endf/b-viii.1" => EndfLibrary::EndfB8_1,
        "jeff3.3" => EndfLibrary::Jeff3_3,
        "jendl5" => EndfLibrary::Jendl5,
        _ => {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "Unknown library '{}'. Use one of: endf8.0, endf8.1, jeff3.3, jendl5",
                library
            )));
        }
    };

    let isotope = Isotope::new(z, a)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid isotope: {}", e)))?;

    let mat_num = match mat {
        Some(m) => m,
        None => mat_number(&isotope).ok_or_else(|| {
            pyo3::exceptions::PyValueError::new_err(format!(
                "MAT number not found for Z={} A={}; provide mat= explicitly",
                z, a
            ))
        })?,
    };

    // Release the GIL for the network I/O (download / cache lookup) and
    // ENDF file parsing.  All types captured by the closure are Send.
    //
    // We tag errors so we can map retrieval failures → PyRuntimeError and
    // parse failures → PyValueError (preserving the pre-GIL-release contract).
    let result: Result<ResonanceData, (bool, String)> = py.detach(move || {
        let retriever = EndfRetriever::new();
        let (_path, contents) = retriever
            .get_endf_file(&isotope, lib, mat_num)
            .map_err(|e| (false, format!("{}", e)))?;

        let data =
            parse_endf_file2(&contents).map_err(|e| (true, format!("ENDF parse error: {}", e)))?;

        Ok(data)
    });

    let data = result.map_err(|(is_parse, msg)| {
        if is_parse {
            pyo3::exceptions::PyValueError::new_err(msg)
        } else {
            pyo3::exceptions::PyRuntimeError::new_err(msg)
        }
    })?;

    // Validate that the parsed ENDF data matches the requested isotope.
    if data.isotope.z() != z || data.isotope.a() != a {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "ENDF data mismatch: requested Z={} A={} but file contains Z={} A={}",
            z,
            a,
            data.isotope.z(),
            data.isotope.a()
        )));
    }

    Ok(PyResonanceData {
        inner: Arc::new(data),
    })
}

/// Load ENDF resonance data from a local file.
///
/// Args:
///     path: Path to an ENDF-format file on disk.
///
/// Returns:
///     ResonanceData parsed from the file.
#[pyfunction]
fn load_endf_file(py: Python<'_>, path: &str) -> PyResult<PyResonanceData> {
    // Release the GIL for the file I/O and ENDF parsing.
    // Tag errors: false = I/O, true = parse.
    let owned_path = path.to_owned();
    let result: Result<ResonanceData, (bool, String)> = py.detach(move || {
        let contents = std::fs::read_to_string(&owned_path)
            .map_err(|e| (false, format!("Cannot read '{}': {}", owned_path, e)))?;

        let data =
            parse_endf_file2(&contents).map_err(|e| (true, format!("ENDF parse error: {}", e)))?;

        Ok(data)
    });

    let data = result.map_err(|(is_parse, msg)| {
        if is_parse {
            pyo3::exceptions::PyValueError::new_err(msg)
        } else {
            pyo3::exceptions::PyIOError::new_err(msg)
        }
    })?;

    Ok(PyResonanceData {
        inner: Arc::new(data),
    })
}

/// Create ResonanceData from parameters (for testing/custom isotopes).
///
/// Args:
///     z: Atomic number.
///     a: Mass number.
///     awr: Atomic weight ratio.
///     scattering_radius: Scattering radius in fm.
///     resonances: List of (energy_eV, j, gn, gg) tuples for L=0.
///     target_spin: Target nuclear spin (default 0.0).
///     l_groups: Optional list of (l_value, [(energy, j, gn, gg), ...]) tuples
///               for multiple L-groups. If provided, the ``resonances`` parameter
///               is ignored.
///     formalism: Resonance formalism to use. Accepted values:
///                - ``None`` or ``"reich_moore"`` (also ``"ReichMoore"``, ``"rm"``,
///                  ``"RM"``, ``"reich-moore"``) — Reich-Moore R-matrix (default).
///                - ``"slbw"`` or ``"SLBW"`` — Single-Level Breit-Wigner.
///
/// Returns:
///     ResonanceData object.
#[pyfunction]
#[pyo3(signature = (z, a, awr, scattering_radius, resonances, target_spin=0.0, l_groups=None, formalism=None))]
fn create_resonance_data(
    z: u32,
    a: u32,
    awr: f64,
    scattering_radius: f64,
    resonances: Vec<(f64, f64, f64, f64)>,
    target_spin: f64,
    l_groups: Option<Vec<(u32, Vec<(f64, f64, f64, f64)>)>>,
    formalism: Option<&str>,
) -> PyResult<PyResonanceData> {
    let res_formalism = match formalism {
        Some("slbw" | "SLBW") => ResonanceFormalism::SLBW,
        Some("reich_moore" | "ReichMoore" | "reich-moore" | "rm" | "RM") | None => {
            ResonanceFormalism::ReichMoore
        }
        Some(other) => {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "Unknown formalism '{other}'. Use 'slbw' or 'reich_moore'."
            )));
        }
    };
    let groups = match l_groups {
        Some(lg) => lg
            .into_iter()
            .map(|(l_val, res_list)| LGroup {
                l: l_val,
                awr,
                apl: 0.0,
                qx: 0.0,
                lrx: 0,
                resonances: res_list
                    .into_iter()
                    .map(|(energy, j, gn, gg)| Resonance {
                        energy,
                        j,
                        gn,
                        gg,
                        gfa: 0.0,
                        gfb: 0.0,
                    })
                    .collect(),
            })
            .collect(),
        None => {
            let res: Vec<Resonance> = resonances
                .into_iter()
                .map(|(energy, j, gn, gg)| Resonance {
                    energy,
                    j,
                    gn,
                    gg,
                    gfa: 0.0,
                    gfb: 0.0,
                })
                .collect();
            vec![LGroup {
                l: 0,
                awr,
                apl: 0.0,
                qx: 0.0,
                lrx: 0,
                resonances: res,
            }]
        }
    };

    Ok(PyResonanceData {
        inner: Arc::new(ResonanceData {
            isotope: Isotope::new(z, a).map_err(|e| {
                pyo3::exceptions::PyValueError::new_err(format!("Invalid isotope: {}", e))
            })?,
            za: z * 1000 + a,
            awr,
            ranges: vec![ResonanceRange {
                energy_low: 1e-5,
                energy_high: 1e6,
                resolved: true,
                formalism: res_formalism,
                target_spin,
                scattering_radius,
                naps: 1,
                l_groups: groups,
                rml: None,
                urr: None,
                ap_table: None,
                r_external: vec![],
            }],
        }),
    })
}

/// Beer-Lambert transmission: T = exp(-thickness * sigma).
///
/// Args:
///     cross_sections: Cross-sections in barns (1D numpy array).
///     thickness: Areal density in atoms/barn.
///
/// Returns:
///     1D numpy array of transmission values.
#[pyfunction]
fn beer_lambert<'py>(
    py: Python<'py>,
    cross_sections: PyReadonlyArray1<f64>,
    thickness: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let xs = cross_sections.as_slice()?;
    let t = transmission::beer_lambert(xs, thickness);
    Ok(PyArray1::from_vec(py, t))
}

/// Validate that an energy grid is finite, positive, and sorted ascending.
/// Empty grids are accepted (callers that need non-empty should use
/// `require_non_empty_energy_grid` instead).
fn validate_energy_grid(e: &[f64]) -> PyResult<()> {
    if e.is_empty() {
        return Ok(());
    }
    if !e[0].is_finite() || e[0] <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "energies must be finite and positive",
        ));
    }
    for i in 1..e.len() {
        if !e[i].is_finite() || e[i] <= 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "energies must be finite and positive",
            ));
        }
        if e[i] <= e[i - 1] {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "energies must be sorted in strictly ascending order",
            ));
        }
    }
    Ok(())
}

/// Validate that an energy grid is **non-empty**, finite, positive, and sorted.
fn require_non_empty_energy_grid(e: &[f64]) -> PyResult<()> {
    if e.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "energies must not be empty",
        ));
    }
    validate_energy_grid(e)
}

/// Build a `ResolutionFunction` from Python arguments.
///
/// Validates mutual exclusivity (Gaussian vs. tabulated) and completeness
/// of Gaussian parameters. Returns `None` when no resolution is requested.
fn build_resolution(
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
) -> PyResult<Option<ResolutionFunction>> {
    let has_gaussian = flight_path_m.is_some() || delta_t_us.is_some() || delta_l_m.is_some();
    if has_gaussian && resolution.is_some() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "Cannot specify both Gaussian resolution parameters and tabulated resolution",
        ));
    }
    let all_gaussian = flight_path_m.is_some() && delta_t_us.is_some() && delta_l_m.is_some();
    if has_gaussian && !all_gaussian {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "Gaussian resolution requires all three parameters: flight_path_m, delta_t_us, and delta_l_m",
        ));
    }
    if let Some(tab) = resolution {
        Ok(Some(ResolutionFunction::Tabulated(tab.inner)))
    } else if let (Some(fp), Some(dt), Some(dl)) = (flight_path_m, delta_t_us, delta_l_m) {
        let de = delta_e_us.unwrap_or(0.0);
        let rp = ResolutionParams::new(fp, dt, dl, de)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(Some(ResolutionFunction::Gaussian(rp)))
    } else {
        Ok(None)
    }
}

/// Apply Free Gas Model (FGM) Doppler broadening to a cross-section array.
///
/// Convolves the input cross-sections with a Gaussian kernel whose width
/// depends on the sample temperature and atomic weight ratio. This is the
/// same broadening applied internally by `forward_model()`, but exposed here
/// so users can broaden individual components (capture, elastic, fission)
/// independently.
///
/// Args:
///     energies: Energy grid in eV (1D numpy array, sorted ascending).
///     cross_sections: Cross-sections in barns (1D numpy array, same length).
///     awr: Atomic weight ratio (target mass / neutron mass).
///     temperature_k: Sample temperature in Kelvin.
///
/// Returns:
///     1D numpy array of Doppler-broadened cross-sections in barns.
///
/// Reference:
///     SAMMY Manual Section III.B.1 (Free-Gas Model of Doppler Broadening).
#[pyfunction]
#[pyo3(signature = (energies, cross_sections, awr, temperature_k))]
fn doppler_broaden<'py>(
    py: Python<'py>,
    energies: PyReadonlyArray1<f64>,
    cross_sections: PyReadonlyArray1<f64>,
    awr: f64,
    temperature_k: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let e = energies.as_slice()?;
    let xs = cross_sections.as_slice()?;

    if e.len() != xs.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "energies length ({}) must match cross_sections length ({})",
            e.len(),
            xs.len(),
        )));
    }
    validate_energy_grid(e)?;

    if temperature_k == 0.0 {
        return Ok(PyArray1::from_vec(py, xs.to_vec()));
    }

    let params = DopplerParams::new(temperature_k, awr).map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("invalid DopplerParams: {e}"))
    })?;

    // Copy numpy slices to owned vectors so we can release the GIL.
    let e_owned = e.to_vec();
    let xs_owned = xs.to_vec();

    // Release the GIL for the Doppler broadening convolution.
    let result = py.detach(move || {
        doppler::doppler_broaden(&e_owned, &xs_owned, &params)
            .map_err(|e| format!("doppler_broaden failed: {e}"))
    });
    let result = result.map_err(|e| pyo3::exceptions::PyValueError::new_err(e))?;
    Ok(PyArray1::from_vec(py, result))
}

/// Apply Gaussian resolution broadening to a cross-section or spectrum array.
///
/// Convolves the input with an energy-dependent Gaussian kernel derived from
/// the instrument's timing uncertainty and flight path length uncertainty.
/// This is the same broadening applied internally by `forward_model()`, but
/// exposed here for independent use on arbitrary arrays.
///
/// Args:
///     energies: Energy grid in eV (1D numpy array, sorted ascending).
///     cross_sections: Values to broaden (1D numpy array, same length).
///     flight_path_m: Flight path length in meters (source to detector).
///     delta_t_us: Total timing uncertainty (1σ Gaussian) in microseconds.
///     delta_l_m: Flight path uncertainty (1σ Gaussian) in meters.
///
/// Returns:
///     1D numpy array of resolution-broadened values.
///
/// Reference:
///     SAMMY Manual Section 3.2 (Resolution Broadening).
#[pyfunction]
#[pyo3(signature = (energies, cross_sections, flight_path_m, delta_t_us, delta_l_m, delta_e_us=0.0))]
fn resolution_broaden<'py>(
    py: Python<'py>,
    energies: PyReadonlyArray1<f64>,
    cross_sections: PyReadonlyArray1<f64>,
    flight_path_m: f64,
    delta_t_us: f64,
    delta_l_m: f64,
    delta_e_us: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let e = energies.as_slice()?;
    let xs = cross_sections.as_slice()?;

    if e.len() != xs.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "energies length ({}) must match cross_sections length ({})",
            e.len(),
            xs.len(),
        )));
    }
    validate_energy_grid(e)?;
    let params = ResolutionParams::new(flight_path_m, delta_t_us, delta_l_m, delta_e_us)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

    if delta_t_us == 0.0 && delta_l_m == 0.0 {
        return Ok(PyArray1::from_vec(py, xs.to_vec()));
    }

    // Copy numpy slices to owned vectors so we can release the GIL.
    let e_owned = e.to_vec();
    let xs_owned = xs.to_vec();

    // Release the GIL for the resolution broadening convolution.
    let result = py.detach(move || resolution::resolution_broaden(&e_owned, &xs_owned, &params));
    let result = result.map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{}", e)))?;
    Ok(PyArray1::from_vec(py, result))
}

/// Load a tabulated resolution function from a VENUS/FTS-format file.
///
/// The file contains reference kernels R(Δt; E_ref) at discrete energies,
/// stored as (TOF_offset_μs, weight) pairs. Kernels are interpolated between
/// reference energies and converted from TOF to energy space during broadening.
///
/// Args:
///     path: Path to the resolution file.
///     flight_path_m: Flight path length in meters (source to detector).
///
/// Returns:
///     TabulatedResolution object for use with ``forward_model()`` or
///     ``fit_spectrum()``.
#[pyfunction]
fn load_resolution(path: &str, flight_path_m: f64) -> PyResult<PyTabulatedResolution> {
    if !flight_path_m.is_finite() || flight_path_m <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "flight_path_m must be finite and positive",
        ));
    }

    let tab = TabulatedResolution::from_file(path, flight_path_m)
        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{}", e)))?;

    Ok(PyTabulatedResolution {
        inner: Arc::new(tab),
    })
}

/// Apply tabulated resolution broadening to a spectrum.
///
/// Convolves the input spectrum with the tabulated instrument resolution
/// function. For each energy point, the kernel is interpolated between
/// reference energies and converted from TOF-offset space to energy space.
///
/// Args:
///     energies: Energy grid in eV (1D numpy array, sorted ascending).
///     spectrum: Values to broaden (1D numpy array, same length).
///     resolution: TabulatedResolution from ``load_resolution()``.
///
/// Returns:
///     1D numpy array of resolution-broadened values.
#[pyfunction]
#[pyo3(name = "apply_resolution")]
fn py_apply_resolution<'py>(
    py: Python<'py>,
    energies: PyReadonlyArray1<f64>,
    spectrum: PyReadonlyArray1<f64>,
    resolution: PyTabulatedResolution,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let e = energies.as_slice()?;
    let s = spectrum.as_slice()?;

    if e.len() != s.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "energies length ({}) must match spectrum length ({})",
            e.len(),
            s.len(),
        )));
    }
    validate_energy_grid(e)?;

    let res_fn = ResolutionFunction::Tabulated(resolution.inner);

    // Copy numpy slices to owned vectors so we can release the GIL.
    let e_owned = e.to_vec();
    let s_owned = s.to_vec();

    // Release the GIL for the tabulated resolution broadening.
    let result = py.detach(move || resolution::apply_resolution(&e_owned, &s_owned, &res_fn));
    let result = result.map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{}", e)))?;
    Ok(PyArray1::from_vec(py, result))
}

/// Run per-pixel fitting across a transmission image stack.
///
/// Routes to either the Levenberg-Marquardt (LM/χ²) fitter or the
/// Poisson/KL sparse fitter depending on the `fitter` argument.
///
/// Args:
///     transmission: 3D array (n_energies, height, width).
///         - fitter='lm':      normalised transmission T ∈ [0,1].
///         - fitter='poisson': raw sample counts (integer-valued floats).
///     uncertainty: 3D array (n_energies, height, width).
///         - fitter='lm':      per-bin transmission uncertainty σ (standard deviation).
///         - fitter='poisson': raw open-beam counts (NOT uncertainty — this parameter is
///                            reused to pass open-beam data; same shape as sample).
///     energies: 1D numpy array of energy values in eV.
///     isotopes: List of ResonanceData objects.
///     temperature_k: Sample temperature in Kelvin (default 300.0).
///     initial_densities: Initial guesses for areal densities (optional).
///     dead_pixels: 2D bool numpy array marking dead pixels (optional).
///     flight_path_m: Flight path for Gaussian resolution (optional).
///     delta_t_us: Timing uncertainty for Gaussian resolution (optional).
///     delta_l_m: Path length uncertainty for Gaussian resolution (optional).
///     resolution: TabulatedResolution for tabulated broadening (optional).
///     max_iter: Maximum iterations per pixel for LM or Poisson fitting (default 100).
///     fitter: 'lm' (default) for Gaussian χ² or 'poisson' for Poisson NLL.
///     roi: [y0, y1, x0, x1] region for Stage-1 nuisance estimation
///         (Poisson path only, default uses full image).
///
/// Returns:
///     SpatialResult (fitter='lm') or SparseResult (fitter='poisson').
#[pyfunction]
#[pyo3(name = "spatial_map", signature = (transmission, uncertainty, energies, isotopes, temperature_k=300.0, initial_densities=None, dead_pixels=None, flight_path_m=None, delta_t_us=None, delta_l_m=None, resolution=None, delta_e_us=None, max_iter=100, fitter="lm", roi=None))]
fn py_spatial_map(
    py: Python<'_>,
    transmission: PyReadonlyArray3<f64>,
    uncertainty: PyReadonlyArray3<f64>,
    energies: PyReadonlyArray1<f64>,
    isotopes: Vec<PyResonanceData>,
    temperature_k: f64,
    initial_densities: Option<Vec<f64>>,
    dead_pixels: Option<PyReadonlyArray2<bool>>,
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
    max_iter: usize,
    fitter: &str,
    roi: Option<[usize; 4]>,
) -> PyResult<Py<PyAny>> {
    let e = energies.as_slice()?;

    if e.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "energies must not be empty",
        ));
    }

    // Validate shapes using cheap PyReadonly views before cloning
    let t_shape = transmission.shape();
    let u_shape = uncertainty.shape();
    if t_shape != u_shape {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "transmission shape {:?} must match uncertainty shape {:?}",
            t_shape, u_shape,
        )));
    }
    if e.len() != t_shape[0] {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "energies length ({}) must match spectral axis ({}) of transmission",
            e.len(),
            t_shape[0],
        )));
    }
    if let Some(ref dead) = dead_pixels {
        let d_shape = dead.shape();
        if d_shape[0] != t_shape[1] || d_shape[1] != t_shape[2] {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "dead_pixels shape ({}, {}) must match spatial dimensions ({}, {}) of transmission",
                d_shape[0], d_shape[1], t_shape[1], t_shape[2],
            )));
        }
    }

    if isotopes.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "isotopes list must not be empty",
        ));
    }

    let n_isotopes = isotopes.len();
    let isotope_names: Vec<String> = isotopes
        .iter()
        .map(|d| {
            let sym = elements::element_symbol(d.inner.isotope.z()).unwrap_or("?");
            format!("{}-{}", sym, d.inner.isotope.a())
        })
        .collect();
    let res_data: Vec<ResonanceData> = isotopes
        .into_iter()
        .map(|d| Arc::unwrap_or_clone(d.inner))
        .collect();

    let init = initial_densities.unwrap_or_else(|| vec![0.001; n_isotopes]);
    if init.len() != n_isotopes {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "initial_densities length ({}) must match isotopes length ({})",
            init.len(),
            n_isotopes,
        )));
    }

    let res_fn = build_resolution(flight_path_m, delta_t_us, delta_l_m, resolution, delta_e_us)?;

    match fitter {
        "lm" => {
            if roi.is_some() {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "roi is only supported with fitter='poisson', not fitter='lm'",
                ));
            }
            let config = FitConfig::new(
                e.to_vec(),
                res_data,
                isotope_names.clone(),
                temperature_k,
                res_fn,
                init,
                LmConfig {
                    max_iter,
                    ..LmConfig::default()
                },
            )
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

            // Clone arrays only after all validation passes.
            // The .to_owned() is necessary because py.detach() releases the GIL and
            // the closure needs 'static data; numpy views borrow GIL-protected memory.
            let trans = transmission.as_array().to_owned();
            let unc = uncertainty.as_array().to_owned();
            let dead = dead_pixels.map(|d| d.as_array().to_owned());

            // Release the GIL for the heavy per-pixel fitting.
            // Pass views into spatial_map to avoid a second deep copy — the function
            // internally transposes into its own owned layout.
            let result = py.detach(move || {
                nereids_pipeline::spatial::spatial_map(
                    trans.view(),
                    unc.view(),
                    &config,
                    dead.as_ref(),
                    None,
                    None,
                )
            });
            let result =
                result.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

            let shape = (
                result.converged_map.shape()[0],
                result.converged_map.shape()[1],
            );
            let py_result = PySpatialResult {
                density_maps: result
                    .density_maps
                    .into_iter()
                    .map(|m| PyArray2::from_owned_array(py, m).unbind())
                    .collect(),
                uncertainty_maps: result
                    .uncertainty_maps
                    .into_iter()
                    .map(|m| PyArray2::from_owned_array(py, m).unbind())
                    .collect(),
                chi_squared_map: PyArray2::from_owned_array(py, result.chi_squared_map).unbind(),
                converged_map: PyArray2::from_owned_array(py, result.converged_map).unbind(),
                n_converged: result.n_converged,
                n_total: result.n_total,
                isotope_names: result.isotope_labels,
                shape,
            };
            Py::new(py, py_result).map(|p| p.into_any())
        }

        "poisson" => {
            // Stage 1: estimate flux + background from open-beam (second arg).
            let sample = transmission.as_array().to_owned();
            let open_beam = uncertainty.as_array().to_owned();
            let dead = dead_pixels.map(|d| d.as_array().to_owned());

            // Validate ROI bounds before feeding into estimate_nuisance.
            let roi_ranges = match roi {
                Some(r) => {
                    let (height, width) = (sample.shape()[1], sample.shape()[2]);
                    if r[0] >= r[1] || r[2] >= r[3] {
                        return Err(pyo3::exceptions::PyValueError::new_err(format!(
                            "roi must be [y0, y1, x0, x1] with y0<y1 and x0<x1, got {:?}",
                            r,
                        )));
                    }
                    if r[0] >= height || r[1] > height || r[2] >= width || r[3] > width {
                        return Err(pyo3::exceptions::PyValueError::new_err(format!(
                            "roi [{},{},{},{}] exceeds image dimensions ({}×{})",
                            r[0], r[1], r[2], r[3], height, width,
                        )));
                    }
                    Some((r[0]..r[1], r[2]..r[3]))
                }
                None => None,
            };
            let sparse_config = SparseConfig::new(
                e.to_vec(),
                res_data,
                isotope_names.clone(),
                temperature_k,
                res_fn,
                init,
                PoissonConfig {
                    max_iter,
                    ..PoissonConfig::default()
                },
            )
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

            // Release the GIL for nuisance estimation + per-pixel Poisson fitting.
            let result = py.detach(move || {
                let nuisance = estimate_nuisance(&open_beam, roi_ranges, dead.as_ref())?;
                sparse_reconstruct(&sample, &nuisance, &sparse_config, dead.as_ref(), None)
            });
            let result =
                result.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

            let shape = (result.nll_map.shape()[0], result.nll_map.shape()[1]);
            let py_result = PySparseResult {
                density_maps: result
                    .density_maps
                    .into_iter()
                    .map(|m| PyArray2::from_owned_array(py, m).unbind())
                    .collect(),
                nll_map: PyArray2::from_owned_array(py, result.nll_map).unbind(),
                converged_map: PyArray2::from_owned_array(py, result.converged_map).unbind(),
                n_converged: result.n_converged,
                n_total: result.n_total,
                isotope_names,
                shape,
            };
            Py::new(py, py_result).map(|p| p.into_any())
        }

        other => Err(pyo3::exceptions::PyValueError::new_err(format!(
            "fitter must be 'lm' or 'poisson', got '{}'",
            other,
        ))),
    }
}

/// Fit a single spectrum averaged over a region of interest.
///
/// Averages the transmission and uncertainty across the specified rectangular
/// region, then fits the resulting high-statistics spectrum. Useful for
/// getting reference densities before per-pixel mapping.
///
/// Args:
///     transmission: 3D numpy array (n_energies, height, width).
///     uncertainty: 3D numpy array (n_energies, height, width).
///     y_range: (start, end) row range for the ROI.
///     x_range: (start, end) column range for the ROI.
///     energies: 1D numpy array of energy values in eV.
///     isotopes: List of ResonanceData objects.
///     temperature_k: Sample temperature in Kelvin (default 300.0, fixed).
///     initial_densities: Initial guesses for areal densities (optional).
///     flight_path_m: Flight path for Gaussian resolution (optional).
///     delta_t_us: Timing uncertainty for Gaussian resolution (optional).
///     delta_l_m: Path length uncertainty for Gaussian resolution (optional).
///     resolution: TabulatedResolution for tabulated broadening (optional).
///     max_iter: Maximum LM iterations (default 100).
///
/// Note:
///     Temperature fitting (``fit_temperature``) is only available via
///     ``fit_spectrum()``.  ``fit_roi()`` always treats temperature as fixed
///     because the ROI-averaged spectrum already has high statistics,
///     making a separate ``fit_spectrum()`` call with ``fit_temperature=True``
///     the preferred workflow.
///
/// Returns:
///     FitResult with densities, uncertainties, and fit quality.
#[pyfunction]
#[pyo3(name = "fit_roi", signature = (transmission, uncertainty, y_range, x_range, energies, isotopes, temperature_k=300.0, initial_densities=None, flight_path_m=None, delta_t_us=None, delta_l_m=None, resolution=None, delta_e_us=None, max_iter=100))]
fn py_fit_roi(
    py: Python<'_>,
    transmission: PyReadonlyArray3<f64>,
    uncertainty: PyReadonlyArray3<f64>,
    y_range: (usize, usize),
    x_range: (usize, usize),
    energies: PyReadonlyArray1<f64>,
    isotopes: Vec<PyResonanceData>,
    temperature_k: f64,
    initial_densities: Option<Vec<f64>>,
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
    max_iter: usize,
) -> PyResult<PyFitResult> {
    let e = energies.as_slice()?;

    if e.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "energies must not be empty",
        ));
    }

    // Validate shapes using cheap PyReadonly views before cloning
    let t_shape = transmission.shape();
    let u_shape = uncertainty.shape();
    if t_shape != u_shape {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "transmission shape {:?} must match uncertainty shape {:?}",
            t_shape, u_shape,
        )));
    }
    if e.len() != t_shape[0] {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "energies length ({}) must match spectral axis ({}) of transmission",
            e.len(),
            t_shape[0],
        )));
    }

    if isotopes.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "isotopes list must not be empty",
        ));
    }
    if y_range.0 >= y_range.1 {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "y_range must be non-empty: start ({}) >= end ({})",
            y_range.0, y_range.1,
        )));
    }
    if x_range.0 >= x_range.1 {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "x_range must be non-empty: start ({}) >= end ({})",
            x_range.0, x_range.1,
        )));
    }
    if y_range.1 > t_shape[1] || x_range.1 > t_shape[2] {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "ROI exceeds image dimensions: y_end={} (max {}), x_end={} (max {})",
            y_range.1, t_shape[1], x_range.1, t_shape[2],
        )));
    }

    let n_isotopes = isotopes.len();
    let isotope_names: Vec<String> = isotopes
        .iter()
        .map(|d| {
            let sym = elements::element_symbol(d.inner.isotope.z()).unwrap_or("?");
            format!("{}-{}", sym, d.inner.isotope.a())
        })
        .collect();
    let res_data: Vec<ResonanceData> = isotopes
        .into_iter()
        .map(|d| Arc::unwrap_or_clone(d.inner))
        .collect();

    let init = initial_densities.unwrap_or_else(|| vec![0.001; n_isotopes]);
    if init.len() != n_isotopes {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "initial_densities length ({}) must match isotopes length ({})",
            init.len(),
            n_isotopes,
        )));
    }

    let res_fn = build_resolution(flight_path_m, delta_t_us, delta_l_m, resolution, delta_e_us)?;

    let config = FitConfig::new(
        e.to_vec(),
        res_data,
        isotope_names,
        temperature_k,
        res_fn,
        init,
        LmConfig {
            max_iter,
            ..LmConfig::default()
        },
    )
    .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

    // Clone arrays only after all validation passes.
    // The .to_owned() is necessary because py.detach() releases the GIL and
    // the closure needs 'static data; numpy views borrow GIL-protected memory.
    let trans = transmission.as_array().to_owned();
    let unc = uncertainty.as_array().to_owned();

    // Release the GIL for the heavy ROI fitting.
    // Pass views into fit_roi to avoid a second deep copy — the function
    // internally slices and transposes into its own owned layout.
    let result = py.detach(move || {
        nereids_pipeline::spatial::fit_roi(
            trans.view(),
            unc.view(),
            y_range.0..y_range.1,
            x_range.0..x_range.1,
            &config,
        )
    });
    let result = result.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

    Ok(PyFitResult {
        densities: result.densities,
        uncertainties: result.uncertainties,
        reduced_chi_squared: result.reduced_chi_squared,
        converged: result.converged,
        iterations: result.iterations,
        temperature_k: None,
        temperature_k_unc: None,
    })
}

/// Load a multi-frame TIFF file into a 3D numpy array.
///
/// Each TIFF frame becomes one slice along the first axis.
/// Data is converted to float64 regardless of the source pixel type.
///
/// Args:
///     path: Path to the multi-frame TIFF file.
///
/// Returns:
///     3D numpy array with shape (n_frames, height, width).
#[pyfunction]
fn load_tiff_stack<'py>(py: Python<'py>, path: &str) -> PyResult<Bound<'py, PyArray3<f64>>> {
    let arr = nereids_io::tiff_stack::load_tiff_stack(std::path::Path::new(path))
        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{}", e)))?;
    Ok(PyArray3::from_owned_array(py, arr))
}

/// Load a folder of single-frame TIFFs into a 3D numpy array.
///
/// Files are sorted lexicographically by name, so they should be named with
/// zero-padded indices (e.g., ``frame_0001.tif``, ``frame_0002.tif``, ...).
///
/// Args:
///     folder: Path to the directory containing TIFF files.
///     pattern: Optional glob pattern matched against each filename (not the
///              full path).  Supports ``*`` and ``?`` wildcards
///              (case-insensitive).  Only files with ``.tif`` or ``.tiff``
///              extensions are ever loaded; the pattern adds an additional
///              filename filter on top of that.
///
/// Returns:
///     3D numpy array with shape (n_frames, height, width), dtype float64.
///
/// Raises:
///     FileNotFoundError: If the folder does not exist or no files match.
///     NotADirectoryError: If the provided path is not a directory.
///     ValueError: If matched frames have inconsistent dimensions.
///     IOError: For TIFF decoding errors or other I/O failures.
#[pyfunction]
#[pyo3(signature = (folder, pattern=None))]
fn load_tiff_folder<'py>(
    py: Python<'py>,
    folder: &str,
    pattern: Option<&str>,
) -> PyResult<Bound<'py, PyArray3<f64>>> {
    let arr = nereids_io::tiff_stack::load_tiff_folder(std::path::Path::new(folder), pattern)
        .map_err(|e| match &e {
            nereids_io::error::IoError::NoMatchingFiles { .. } => {
                pyo3::exceptions::PyFileNotFoundError::new_err(format!("{}", e))
            }
            nereids_io::error::IoError::NotADirectory(_) => {
                pyo3::exceptions::PyNotADirectoryError::new_err(format!("{}", e))
            }
            nereids_io::error::IoError::FileNotFound(_, source)
                if source.kind() == std::io::ErrorKind::NotFound =>
            {
                pyo3::exceptions::PyFileNotFoundError::new_err(format!("{}", e))
            }
            nereids_io::error::IoError::DimensionMismatch { .. } => {
                pyo3::exceptions::PyValueError::new_err(format!("{}", e))
            }
            _ => pyo3::exceptions::PyIOError::new_err(format!("{}", e)),
        })?;
    Ok(PyArray3::from_owned_array(py, arr))
}

/// Normalize raw sample and open-beam data to transmission.
///
/// Computes T = (C_sample / C_ob) × (PC_ob / PC_sample) with Poisson
/// uncertainty propagation.
///
/// Args:
///     sample: 3D numpy array of raw sample counts (n_tof, height, width).
///     open_beam: 3D numpy array of open-beam counts (same shape).
///     pc_sample: Proton charge for the sample measurement.
///     pc_ob: Proton charge for the open-beam measurement.
///     dark_current: Optional 2D numpy array (height, width) to subtract.
///
/// Returns:
///     Tuple of (transmission, uncertainty) as 3D numpy arrays.
#[pyfunction]
#[pyo3(signature = (sample, open_beam, pc_sample, pc_ob, dark_current=None))]
fn normalize<'py>(
    py: Python<'py>,
    sample: PyReadonlyArray3<f64>,
    open_beam: PyReadonlyArray3<f64>,
    pc_sample: f64,
    pc_ob: f64,
    dark_current: Option<PyReadonlyArray2<f64>>,
) -> PyResult<(Bound<'py, PyArray3<f64>>, Bound<'py, PyArray3<f64>>)> {
    // Validate shapes using cheap PyReadonly views before cloning
    let s_shape = sample.shape();
    let ob_shape = open_beam.shape();
    if s_shape != ob_shape {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "sample shape {:?} must match open_beam shape {:?}",
            s_shape, ob_shape,
        )));
    }

    if let Some(ref dc_arr) = dark_current {
        let dc_shape = dc_arr.shape();
        if dc_shape[0] != s_shape[1] || dc_shape[1] != s_shape[2] {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "dark_current shape ({}, {}) must match spatial dimensions ({}, {}) of sample",
                dc_shape[0], dc_shape[1], s_shape[1], s_shape[2],
            )));
        }
    }

    // Clone arrays only after all validation passes
    let s = sample.as_array().to_owned();
    let ob = open_beam.as_array().to_owned();
    let dc = dark_current.map(|d| d.as_array().to_owned());

    let params = NormalizationParams {
        proton_charge_sample: pc_sample,
        proton_charge_ob: pc_ob,
    };

    // Release the GIL for the normalization computation.
    let result = py.detach(move || norm::normalize(&s, &ob, &params, dc.as_ref()));
    let result = result.map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{}", e)))?;

    Ok((
        PyArray3::from_owned_array(py, result.transmission),
        PyArray3::from_owned_array(py, result.uncertainty),
    ))
}

/// Convert TOF bin edges to energy bin centers.
///
/// Returns the geometric mean of adjacent energy bin edges (ascending order).
/// This is the standard energy grid for neutron resonance analysis.
///
/// Args:
///     tof_edges: 1D numpy array of TOF bin edges in microseconds (ascending).
///     flight_path_m: Total flight path in meters.
///     delay_us: Electronic/moderator delay in microseconds (default 0.0).
///
/// Returns:
///     1D numpy array of energy bin centers in eV (ascending).
///     Length = len(tof_edges) - 1.
#[pyfunction]
#[pyo3(signature = (tof_edges, flight_path_m, delay_us=0.0))]
fn tof_to_energy_centers<'py>(
    py: Python<'py>,
    tof_edges: PyReadonlyArray1<f64>,
    flight_path_m: f64,
    delay_us: f64,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let edges = tof_edges.as_slice()?;
    let params = BeamlineParams {
        flight_path_m,
        delay_us,
    };

    let centers = nereids_io::tof::tof_edges_to_energy_centers(edges, &params)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{}", e)))?;

    Ok(PyArray1::from_owned_array(py, centers))
}

// ── Element / isotope utilities ──────────────────────────────────────

/// Get the element symbol for a given atomic number Z.
///
/// Args:
///     z: Atomic number (e.g. 92 for uranium).
///
/// Returns:
///     Element symbol (e.g. "U"), or None if Z is out of range.
#[pyfunction]
#[pyo3(name = "element_symbol")]
fn py_element_symbol(z: u32) -> Option<String> {
    elements::element_symbol(z).map(|s| s.to_string())
}

/// Get the element name for a given atomic number Z.
///
/// Args:
///     z: Atomic number (e.g. 92 for uranium).
///
/// Returns:
///     Element name (e.g. "Uranium"), or None if Z is out of range.
#[pyfunction]
#[pyo3(name = "element_name")]
fn py_element_name(z: u32) -> Option<String> {
    elements::element_name(z).map(|s| s.to_string())
}

/// Parse an isotope string like "U-238" into (Z, A).
///
/// Args:
///     s: Isotope string in "Symbol-A" format (e.g. "U-238", "Fe-56").
///
/// Returns:
///     Tuple (z, a) or None if the string cannot be parsed.
#[pyfunction]
#[pyo3(name = "parse_isotope_str")]
fn py_parse_isotope_str(s: &str) -> Option<(u32, u32)> {
    elements::parse_isotope_str(s).map(|iso| (iso.z(), iso.a()))
}

/// Get the natural isotopic abundance for a specific isotope.
///
/// Args:
///     z: Atomic number.
///     a: Mass number.
///
/// Returns:
///     Abundance as a fraction (0.0 to 1.0), or None for synthetic isotopes.
#[pyfunction]
#[pyo3(name = "natural_abundance")]
fn py_natural_abundance(z: u32, a: u32) -> Option<f64> {
    Isotope::new(z, a)
        .ok()
        .and_then(|iso| elements::natural_abundance(&iso))
}

/// Get all naturally occurring isotopes for an element.
///
/// Args:
///     z: Atomic number (e.g. 74 for tungsten).
///
/// Returns:
///     List of ((z, a), abundance) tuples for all stable isotopes.
#[pyfunction]
#[pyo3(name = "natural_isotopes")]
fn py_natural_isotopes(z: u32) -> Vec<((u32, u32), f64)> {
    elements::natural_isotopes(z)
        .into_iter()
        .map(|(iso, frac)| ((iso.z(), iso.a()), frac))
        .collect()
}

/// Result of a trace-detectability analysis.
///
/// Returned by ``trace_detectability()`` and ``trace_detectability_survey()``.
/// Contains the peak SNR, the energy at which peak contrast occurs, and the
/// full |ΔT| spectrum for plotting.
#[pyclass(name = "TraceDetectabilityReport")]
struct PyTraceDetectabilityReport {
    inner: detectability::TraceDetectabilityReport,
}

#[pymethods]
impl PyTraceDetectabilityReport {
    /// Peak |ΔT| per ppm concentration at the most sensitive energy.
    #[getter]
    fn peak_delta_t_per_ppm(&self) -> f64 {
        self.inner.peak_delta_t_per_ppm
    }

    /// Energy at which peak contrast occurs (eV).
    #[getter]
    fn peak_energy_ev(&self) -> f64 {
        self.inner.peak_energy_ev
    }

    /// Estimated peak SNR at the given concentration and I₀.
    #[getter]
    fn peak_snr(&self) -> f64 {
        self.inner.peak_snr
    }

    /// Whether the combination is detectable (SNR > threshold).
    #[getter]
    fn detectable(&self) -> bool {
        self.inner.detectable
    }

    /// Energy-resolved |ΔT| spectrum for the given concentration.
    #[getter]
    fn delta_t_spectrum<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_vec(py, self.inner.delta_t_spectrum.clone())
    }

    /// Energies used (eV).
    #[getter]
    fn energies<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_vec(py, self.inner.energies.clone())
    }

    /// Fraction of energy bins where the matrix baseline is opaque (T < 1e-15).
    #[getter]
    fn opaque_fraction(&self) -> f64 {
        self.inner.opaque_fraction
    }

    fn __repr__(&self) -> String {
        format!(
            "TraceDetectabilityReport(detectable={}, peak_snr={:.2}, peak_energy_ev={:.2}, opaque_fraction={:.2})",
            self.inner.detectable,
            self.inner.peak_snr,
            self.inner.peak_energy_ev,
            self.inner.opaque_fraction,
        )
    }
}

/// Compute trace-detectability for a matrix + trace isotope pair.
///
/// Determines whether a trace isotope is detectable at a given concentration
/// (in ppm) within a matrix, by computing the peak spectral SNR over the
/// supplied energy window.
///
/// Resolution broadening can be applied via either Gaussian parameters
/// (``flight_path_m``, ``delta_t_us``, ``delta_l_m``) or a tabulated
/// resolution function (``resolution``). Providing both is an error.
///
/// Args:
///     matrix: ResonanceData for the matrix isotope.
///     matrix_density: Matrix areal density in atoms/barn.
///     trace: ResonanceData for the trace isotope.
///     trace_ppm: Trace concentration in ppm by atom.
///     energies: Energy grid in eV (1D numpy array, sorted ascending).
///     i0: Expected counts per energy bin (for Poisson noise estimate).
///     temperature_k: Sample temperature in Kelvin (default 293.6).
///     flight_path_m: Flight path for Gaussian resolution (optional).
///     delta_t_us: Timing uncertainty for Gaussian resolution (optional).
///     delta_l_m: Path length uncertainty for Gaussian resolution (optional).
///     resolution: TabulatedResolution for tabulated broadening (optional).
///     snr_threshold: Detection threshold in σ (default 3.0).
///
/// Returns:
///     TraceDetectabilityReport with peak SNR, peak energy, and |ΔT| spectrum.
#[pyfunction]
#[pyo3(name = "trace_detectability", signature = (matrix, matrix_density, trace, trace_ppm, energies, i0, temperature_k=293.6, flight_path_m=None, delta_t_us=None, delta_l_m=None, resolution=None, delta_e_us=None, snr_threshold=3.0))]
fn py_trace_detectability(
    py: Python<'_>,
    matrix: &PyResonanceData,
    matrix_density: f64,
    trace: &PyResonanceData,
    trace_ppm: f64,
    energies: PyReadonlyArray1<f64>,
    i0: f64,
    temperature_k: f64,
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
    snr_threshold: f64,
) -> PyResult<PyTraceDetectabilityReport> {
    let e = energies.as_slice()?;
    require_non_empty_energy_grid(e)?;

    if matrix_density <= 0.0 || !matrix_density.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "matrix_density must be finite and positive",
        ));
    }
    if trace_ppm < 0.0 || !trace_ppm.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "trace_ppm must be finite and non-negative",
        ));
    }
    if i0 <= 0.0 || !i0.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "i0 must be finite and positive",
        ));
    }
    if !temperature_k.is_finite() || temperature_k < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temperature_k must be finite and non-negative",
        ));
    }
    if snr_threshold < 0.0 || !snr_threshold.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "snr_threshold must be finite and non-negative",
        ));
    }

    let res_fn = build_resolution(flight_path_m, delta_t_us, delta_l_m, resolution, delta_e_us)?;

    // Clone data to owned types so we can release the GIL.
    let e_owned = e.to_vec();
    let matrix_data = matrix.inner.clone();
    let trace_data = trace.inner.clone();

    // Release the GIL for the detectability computation.
    // Wrap single matrix in a vec — Rust API supports multi-matrix but Python
    // API preserves backward compatibility with a single matrix argument.
    let report = py.detach(move || {
        let matrix_isotopes = vec![(Arc::unwrap_or_clone(matrix_data), matrix_density)];
        let config = detectability::TraceDetectabilityConfig {
            matrix_isotopes: &matrix_isotopes,
            energies: &e_owned,
            i0,
            temperature_k,
            resolution: res_fn.as_ref(),
            snr_threshold,
        };
        detectability::trace_detectability(&config, &trace_data, trace_ppm)
            .map_err(|e| format!("trace_detectability failed: {e}"))
    });
    let report = report.map_err(|e| pyo3::exceptions::PyValueError::new_err(e))?;

    Ok(PyTraceDetectabilityReport { inner: report })
}

/// Survey multiple trace candidates against a single matrix.
///
/// Parallelises over candidates with rayon. Returns a list of
/// ``(isotope_name, TraceDetectabilityReport)`` tuples sorted by
/// ``peak_snr`` descending.
///
/// Resolution broadening can be applied via either Gaussian parameters
/// (``flight_path_m``, ``delta_t_us``, ``delta_l_m``) or a tabulated
/// resolution function (``resolution``). Providing both is an error.
///
/// Args:
///     matrix: ResonanceData for the matrix isotope.
///     matrix_density: Matrix areal density in atoms/barn.
///     trace_candidates: List of ResonanceData for candidate trace isotopes.
///     trace_ppm: Trace concentration in ppm by atom.
///     energies: Energy grid in eV (1D numpy array, sorted ascending).
///     i0: Expected counts per energy bin (for Poisson noise estimate).
///     temperature_k: Sample temperature in Kelvin (default 293.6).
///     flight_path_m: Flight path for Gaussian resolution (optional).
///     delta_t_us: Timing uncertainty for Gaussian resolution (optional).
///     delta_l_m: Path length uncertainty for Gaussian resolution (optional).
///     resolution: TabulatedResolution for tabulated broadening (optional).
///     snr_threshold: Detection threshold in σ (default 3.0).
///
/// Returns:
///     List of (isotope_name, TraceDetectabilityReport) sorted by peak_snr descending.
#[pyfunction]
#[pyo3(name = "trace_detectability_survey", signature = (matrix, matrix_density, trace_candidates, trace_ppm, energies, i0, temperature_k=293.6, flight_path_m=None, delta_t_us=None, delta_l_m=None, resolution=None, delta_e_us=None, snr_threshold=3.0))]
fn py_trace_detectability_survey(
    py: Python<'_>,
    matrix: &PyResonanceData,
    matrix_density: f64,
    trace_candidates: Vec<PyResonanceData>,
    trace_ppm: f64,
    energies: PyReadonlyArray1<f64>,
    i0: f64,
    temperature_k: f64,
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
    snr_threshold: f64,
) -> PyResult<Vec<(String, PyTraceDetectabilityReport)>> {
    let e = energies.as_slice()?;
    require_non_empty_energy_grid(e)?;

    if matrix_density <= 0.0 || !matrix_density.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "matrix_density must be finite and positive",
        ));
    }
    if trace_ppm < 0.0 || !trace_ppm.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "trace_ppm must be finite and non-negative",
        ));
    }
    if i0 <= 0.0 || !i0.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "i0 must be finite and positive",
        ));
    }
    if !temperature_k.is_finite() || temperature_k < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temperature_k must be finite and non-negative",
        ));
    }
    if snr_threshold < 0.0 || !snr_threshold.is_finite() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "snr_threshold must be finite and non-negative",
        ));
    }
    if trace_candidates.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "trace_candidates must not be empty",
        ));
    }

    let res_fn = build_resolution(flight_path_m, delta_t_us, delta_l_m, resolution, delta_e_us)?;

    let candidates: Vec<ResonanceData> = trace_candidates
        .into_iter()
        .map(|d| Arc::unwrap_or_clone(d.inner))
        .collect();

    // Clone data to owned types so we can release the GIL.
    let e_owned = e.to_vec();
    let matrix_data = matrix.inner.clone();

    // Release the GIL for the parallelised detectability survey.
    // Wrap single matrix in a vec — Rust API supports multi-matrix but Python
    // API preserves backward compatibility with a single matrix argument.
    let results = py.detach(move || {
        let matrix_isotopes = vec![(Arc::unwrap_or_clone(matrix_data), matrix_density)];
        let config = detectability::TraceDetectabilityConfig {
            matrix_isotopes: &matrix_isotopes,
            energies: &e_owned,
            i0,
            temperature_k,
            resolution: res_fn.as_ref(),
            snr_threshold,
        };
        detectability::trace_detectability_survey(&config, &candidates, trace_ppm)
            .map_err(|e| format!("trace_detectability_survey failed: {e}"))
    });
    let results = results.map_err(|e| pyo3::exceptions::PyValueError::new_err(e))?;

    Ok(results
        .into_iter()
        .map(|(name, report)| (name, PyTraceDetectabilityReport { inner: report }))
        .collect())
}

/// Precompute Doppler- and resolution-broadened total cross-sections.
///
/// Returns one broadened total cross-section array per isotope.  This is the
/// expensive physics step (Doppler FGM + resolution convolution); calling it
/// once and caching the result avoids redundant computation when the same
/// isotopes and energy grid are reused across many fits or forward-model
/// evaluations.
///
/// Resolution broadening can be applied via either Gaussian parameters
/// (``flight_path_m``, ``delta_t_us``, ``delta_l_m``) or a tabulated
/// resolution function (``resolution``). Providing both is an error.
///
/// Args:
///     energies: Energy grid in eV (1D numpy array, sorted ascending).
///     isotopes: List of ResonanceData objects.
///     temperature_k: Sample temperature in Kelvin (default 0.0).
///     flight_path_m: Flight path in meters for Gaussian resolution (optional).
///     delta_t_us: Timing uncertainty in microseconds (optional).
///     delta_l_m: Path length uncertainty in meters (optional).
///     resolution: TabulatedResolution from ``load_resolution()`` (optional).
///
/// Returns:
///     List of 1D numpy arrays (one per isotope), each containing the broadened
///     total cross-section in barns on the supplied energy grid.
#[pyfunction]
#[pyo3(signature = (energies, isotopes, temperature_k=0.0, flight_path_m=None, delta_t_us=None, delta_l_m=None, resolution=None, delta_e_us=None))]
fn precompute_cross_sections<'py>(
    py: Python<'py>,
    energies: PyReadonlyArray1<f64>,
    isotopes: Vec<PyResonanceData>,
    temperature_k: f64,
    flight_path_m: Option<f64>,
    delta_t_us: Option<f64>,
    delta_l_m: Option<f64>,
    resolution: Option<PyTabulatedResolution>,
    delta_e_us: Option<f64>,
) -> PyResult<Vec<Bound<'py, PyArray1<f64>>>> {
    let e = energies.as_slice()?;
    require_non_empty_energy_grid(e)?;

    if isotopes.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "isotopes list must not be empty",
        ));
    }
    if !temperature_k.is_finite() || temperature_k < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "temperature_k must be finite and non-negative",
        ));
    }

    let res_data: Vec<ResonanceData> = isotopes
        .into_iter()
        .map(|d| Arc::unwrap_or_clone(d.inner))
        .collect();
    let res_fn = build_resolution(flight_path_m, delta_t_us, delta_l_m, resolution, delta_e_us)?;
    let instrument = res_fn.map(|r| InstrumentParams { resolution: r });

    // Copy numpy slice to owned Vec so we can release the GIL.
    let e_owned = e.to_vec();

    // Release the GIL for the heavy Doppler + resolution broadening.
    let xs = py.detach(move || {
        transmission::broadened_cross_sections(
            &e_owned,
            &res_data,
            temperature_k,
            instrument.as_ref(),
            None,
        )
    });

    // GIL is re-acquired after detach returns — use `py` directly.
    let xs = xs.map_err(|e| {
        pyo3::exceptions::PyRuntimeError::new_err(format!("broadened_cross_sections failed: {}", e))
    })?;

    Ok(xs.into_iter().map(|v| PyArray1::from_vec(py, v)).collect())
}

/// Detect dead pixels in a 3D image stack.
///
/// A pixel is marked as "dead" when all its counts across the spectral/TOF
/// axis are exactly zero.  The returned mask can be passed directly to
/// ``spatial_map(dead_pixels=...)``.
///
/// Args:
///     data: 3D numpy array with shape ``(n_frames, height, width)``.
///         Typically an open-beam stack or raw sample stack.
///
/// Returns:
///     2D boolean numpy array with shape ``(height, width)``.
///     ``True`` marks a dead pixel.
#[pyfunction]
fn detect_dead_pixels<'py>(
    py: Python<'py>,
    data: PyReadonlyArray3<f64>,
) -> PyResult<Bound<'py, PyArray2<bool>>> {
    let arr = data.as_array().to_owned();
    let mask = py.detach(move || norm::detect_dead_pixels(&arr));
    Ok(PyArray2::from_owned_array(py, mask))
}

/// NEREIDS Python module.
#[pymodule]
fn nereids(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyResonanceData>()?;
    m.add_class::<PyFitResult>()?;
    m.add_class::<PyTabulatedResolution>()?;
    m.add_class::<PySpatialResult>()?;
    m.add_class::<PySparseResult>()?;
    m.add_class::<PyTraceDetectabilityReport>()?;
    m.add_function(wrap_pyfunction!(cross_sections, m)?)?;
    m.add_function(wrap_pyfunction!(forward_model, m)?)?;
    m.add_function(wrap_pyfunction!(fit_spectrum, m)?)?;
    m.add_function(wrap_pyfunction!(tof_to_energy, m)?)?;
    m.add_function(wrap_pyfunction!(energy_to_tof, m)?)?;
    m.add_function(wrap_pyfunction!(load_endf, m)?)?;
    m.add_function(wrap_pyfunction!(load_endf_file, m)?)?;
    m.add_function(wrap_pyfunction!(create_resonance_data, m)?)?;
    m.add_function(wrap_pyfunction!(beer_lambert, m)?)?;
    m.add_function(wrap_pyfunction!(doppler_broaden, m)?)?;
    m.add_function(wrap_pyfunction!(resolution_broaden, m)?)?;
    m.add_function(wrap_pyfunction!(load_resolution, m)?)?;
    m.add_function(wrap_pyfunction!(py_apply_resolution, m)?)?;
    m.add_function(wrap_pyfunction!(py_spatial_map, m)?)?;
    m.add_function(wrap_pyfunction!(py_fit_roi, m)?)?;
    m.add_function(wrap_pyfunction!(load_tiff_stack, m)?)?;
    m.add_function(wrap_pyfunction!(load_tiff_folder, m)?)?;
    m.add_function(wrap_pyfunction!(normalize, m)?)?;
    m.add_function(wrap_pyfunction!(tof_to_energy_centers, m)?)?;
    m.add_function(wrap_pyfunction!(py_element_symbol, m)?)?;
    m.add_function(wrap_pyfunction!(py_element_name, m)?)?;
    m.add_function(wrap_pyfunction!(py_parse_isotope_str, m)?)?;
    m.add_function(wrap_pyfunction!(py_natural_abundance, m)?)?;
    m.add_function(wrap_pyfunction!(py_natural_isotopes, m)?)?;
    m.add_function(wrap_pyfunction!(py_trace_detectability, m)?)?;
    m.add_function(wrap_pyfunction!(py_trace_detectability_survey, m)?)?;
    m.add_function(wrap_pyfunction!(precompute_cross_sections, m)?)?;
    m.add_function(wrap_pyfunction!(detect_dead_pixels, m)?)?;
    Ok(())
}
