//! Trace-detectability analysis for neutron resonance imaging.
//!
//! Answers the pre-experiment question: *"Can I detect X ppm of isotope B
//! in a matrix of isotope(s) A across a given energy window?"*
//!
//! ## Core concept
//!
//! For a given matrix + trace isotope pair, compute the **peak spectral SNR**
//! as a function of trace concentration:
//!
//!   SNR_peak(c) = max_E |ΔT(E, c)| / σ_noise
//!
//! where ΔT is the *signed* transmission difference:
//!
//!   ΔT(E, c) = T(E, n_matrix, 0) − T(E, n_matrix, n_trace = c·n_matrix)
//!
//! and σ_noise ≈ 1/√I₀ (off-resonance Poisson approximation).
//!
//! The stored `delta_t_spectrum` and all derived metrics (peak_delta_t_per_ppm,
//! peak_snr) use **|ΔT|**, discarding the sign.
//!
//! ## Reference
//!
//! Motivated by the observation that Fe-56 + Mn-55 have no resolved resonances
//! in 1–50 eV, while W-182 + Hf-178 give strong contrast in the same window.
//! VENUS can resolve up to ~1 keV, so many more pairs become accessible at
//! higher energies (e.g., Mn-55 resonance at ~337 eV).

use nereids_core::elements;
use nereids_endf::resonance::ResonanceData;
use nereids_physics::resolution::ResolutionFunction;
use nereids_physics::transmission::{self, InstrumentParams, SampleParams, TransmissionError};
use rayon::prelude::*;

use crate::error::PipelineError;

/// Transmission threshold below which a bin is considered opaque.
///
/// At `T < 1e-15`, `exp(-n·σ)` is effectively zero in f64 arithmetic.
/// Any ΔT computation between two opaque transmissions yields 0.0 due to
/// floating-point underflow, making trace detection impossible regardless
/// of the actual physics.
const OPAQUE_THRESHOLD: f64 = 1e-15;

/// Result of a trace-detectability analysis for a single matrix+trace pair.
#[derive(Debug, Clone)]
pub struct TraceDetectabilityReport {
    /// Peak |ΔT| per ppm concentration at the most sensitive energy.
    ///
    /// Linearized sensitivity metric: `peak_delta_t / trace_ppm`.
    /// Accurate in the dilute-trace limit; at high concentrations the
    /// actual per-ppm sensitivity is smaller due to Beer-Lambert
    /// saturation (T = exp(−n·σ) is sub-linear in n).
    pub peak_delta_t_per_ppm: f64,
    /// Energy at which peak contrast occurs (eV).
    pub peak_energy_ev: f64,
    /// Estimated peak SNR at the given concentration and I₀.
    pub peak_snr: f64,
    /// Whether the combination is detectable (SNR > threshold).
    pub detectable: bool,
    /// Energy-resolved |ΔT| spectrum for the given concentration.
    pub delta_t_spectrum: Vec<f64>,
    /// Energies used (eV).
    pub energies: Vec<f64>,
    /// Fraction of energy bins where the matrix baseline transmission is
    /// below `OPAQUE_THRESHOLD` (effectively zero due to underflow).
    ///
    /// When this is close to 1.0, the matrix is essentially opaque across
    /// the entire energy range, and trace detection is physically impossible
    /// regardless of concentration or I₀.
    pub opaque_fraction: f64,
}

/// Configuration for a trace-detectability analysis.
pub struct TraceDetectabilityConfig<'a> {
    /// Matrix isotopes with their areal densities in atoms/barn.
    ///
    /// Supports multi-isotope matrices (e.g., an alloy). The total matrix
    /// density (sum of all densities) is used to convert trace ppm to
    /// areal density: `trace_density = trace_ppm × 1e-6 × total_matrix_density`.
    pub matrix_isotopes: &'a [(ResonanceData, f64)],
    /// Energy grid in eV (sorted ascending).
    pub energies: &'a [f64],
    /// Expected counts per bin (for Poisson noise estimate).
    pub i0: f64,
    /// Sample temperature in Kelvin.
    pub temperature_k: f64,
    /// Optional resolution broadening function.
    pub resolution: Option<&'a ResolutionFunction>,
    /// Detection threshold (3.0 = standard 3σ detection limit).
    pub snr_threshold: f64,
}

/// Build an `InstrumentParams` from the config's optional resolution function.
fn build_instrument(config: &TraceDetectabilityConfig) -> Option<InstrumentParams> {
    config.resolution.map(|r| InstrumentParams {
        resolution: r.clone(),
    })
}

/// Total matrix areal density (sum of all matrix isotope densities).
fn total_matrix_density(config: &TraceDetectabilityConfig) -> f64 {
    config.matrix_isotopes.iter().map(|(_, d)| d).sum()
}

/// Compute the matrix-only baseline transmission.
///
/// NOTE: `SampleParams` owns its `ResonanceData`, so we clone here.
/// For typical ENDF isotopes the data is small (a few KB) and the clone
/// cost is negligible compared to the cross-section evaluation.  If
/// `SampleParams` gains `Arc` support in the future, these clones can
/// be eliminated.
fn matrix_baseline(
    config: &TraceDetectabilityConfig,
    instrument: Option<&InstrumentParams>,
) -> Result<Vec<f64>, TransmissionError> {
    let isotopes: Vec<_> = config.matrix_isotopes.to_vec();
    let sample_matrix = SampleParams::new(config.temperature_k, isotopes)
        .map_err(|e| TransmissionError::InputMismatch(e.to_string()))?;
    transmission::forward_model(config.energies, &sample_matrix, instrument)
}

/// Build a report from a precomputed matrix-only baseline and a trace isotope.
fn report_from_baseline(
    config: &TraceDetectabilityConfig,
    instrument: Option<&InstrumentParams>,
    t_matrix: &[f64],
    trace: &ResonanceData,
    trace_ppm: f64,
) -> Result<TraceDetectabilityReport, TransmissionError> {
    // T_combined: transmission through matrix + trace
    let trace_density = trace_ppm * 1e-6 * total_matrix_density(config);
    let mut isotopes: Vec<_> = config.matrix_isotopes.to_vec();
    isotopes.push((trace.clone(), trace_density));
    let sample_combined = SampleParams::new(config.temperature_k, isotopes)
        .map_err(|e| TransmissionError::InputMismatch(e.to_string()))?;
    let t_combined = transmission::forward_model(config.energies, &sample_combined, instrument)?;

    // Opaque fraction: bins where matrix baseline underflows to ~0.
    let n_opaque = t_matrix.iter().filter(|&&t| t < OPAQUE_THRESHOLD).count();
    let opaque_fraction = if t_matrix.is_empty() {
        0.0
    } else {
        n_opaque as f64 / t_matrix.len() as f64
    };

    // |ΔT| spectrum — absolute difference, sign discarded (see module docs).
    let delta_t_spectrum: Vec<f64> = t_matrix
        .iter()
        .zip(t_combined.iter())
        .map(|(&tm, &tc)| (tm - tc).abs())
        .collect();

    // Find peak |ΔT| and corresponding energy
    let (peak_idx, &peak_delta_t) = delta_t_spectrum
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .unwrap_or((0, &0.0));

    let peak_energy_ev = if config.energies.is_empty() {
        0.0
    } else {
        config.energies[peak_idx]
    };

    // P-8: SNR uses energy-dependent noise σ(E) = √(T_matrix(E) / I₀).
    //
    // Previously used the off-resonance approximation σ ≈ 1/√I₀ (T≈1),
    // which underestimates SNR at resonance dips where T<<1 and noise
    // is correspondingly smaller.  The energy-dependent formula gives
    // the correct Poisson counting noise at each bin.
    //
    // Peak SNR = max_E |ΔT(E)| / σ(E)  where σ(E) = √(T_matrix(E) / I₀)
    let peak_snr = if config.i0 > 0.0 {
        delta_t_spectrum
            .iter()
            .zip(t_matrix.iter())
            .map(|(&dt, &tm)| {
                let sigma_e = (tm.max(OPAQUE_THRESHOLD) / config.i0).sqrt();
                dt / sigma_e
            })
            .fold(0.0f64, f64::max)
    } else {
        0.0
    };

    // Per-ppm normalisation
    let peak_delta_t_per_ppm = if trace_ppm > 0.0 {
        peak_delta_t / trace_ppm
    } else {
        0.0
    };

    Ok(TraceDetectabilityReport {
        peak_delta_t_per_ppm,
        peak_energy_ev,
        peak_snr,
        detectable: peak_snr > config.snr_threshold,
        delta_t_spectrum,
        energies: config.energies.to_vec(),
        opaque_fraction,
    })
}

/// Compute trace-detectability for a single matrix+trace isotope pair.
///
/// Two `forward_model` calls (with/without trace) + argmax over the
/// difference spectrum. Resolution broadening is included when provided,
/// so the SNR reflects realistic peak spreading at VENUS.
///
/// # Arguments
/// * `config`    — Shared analysis parameters (matrix isotopes, energies, I₀, etc.).
/// * `trace`     — Resonance data for the trace isotope.
/// * `trace_ppm` — Trace concentration in ppm by atom.
///
/// # Preconditions
/// * `config.energies` must be non-empty and sorted ascending.
/// * `config.i0` must be positive (used as `1/√I₀` for noise estimate).
/// * `config.matrix_isotopes` must be non-empty with positive densities.
/// * `config.snr_threshold` must be non-negative.
/// * `trace_ppm` must be non-negative.
///
/// The Python bindings validate all of these; Rust callers are responsible
/// for ensuring valid inputs.
///
/// # Returns
/// [`TraceDetectabilityReport`] with peak SNR, peak energy, detectability
/// flag, and the full |ΔT| spectrum.
pub fn trace_detectability(
    config: &TraceDetectabilityConfig,
    trace: &ResonanceData,
    trace_ppm: f64,
) -> Result<TraceDetectabilityReport, PipelineError> {
    let instrument = build_instrument(config);
    let t_matrix = matrix_baseline(config, instrument.as_ref())?;
    Ok(report_from_baseline(
        config,
        instrument.as_ref(),
        &t_matrix,
        trace,
        trace_ppm,
    )?)
}

/// Survey multiple trace candidates against a matrix (single or multi-isotope).
///
/// The matrix-only baseline transmission and instrument resolution are
/// computed once and reused for all candidates. Each candidate is then
/// evaluated in parallel with rayon.
///
/// # Preconditions
/// Same as [`trace_detectability`], plus `trace_candidates` must be
/// non-empty.
///
/// # Returns
/// Vec of (isotope_name, report) sorted by `peak_snr` descending.
pub fn trace_detectability_survey(
    config: &TraceDetectabilityConfig,
    trace_candidates: &[ResonanceData],
    trace_ppm: f64,
) -> Result<Vec<(String, TraceDetectabilityReport)>, PipelineError> {
    // Build instrument and matrix baseline once for all candidates.
    let instrument = build_instrument(config);
    let t_matrix = matrix_baseline(config, instrument.as_ref())?;

    let mut results: Vec<(String, TraceDetectabilityReport)> = trace_candidates
        .par_iter()
        .map(|trace| {
            let name = elements::element_symbol(trace.isotope.z())
                .map(|sym| format!("{}-{}", sym, trace.isotope.a()))
                .unwrap_or_else(|| format!("Z{}-{}", trace.isotope.z(), trace.isotope.a()));

            let report =
                report_from_baseline(config, instrument.as_ref(), &t_matrix, trace, trace_ppm)?;

            Ok((name, report))
        })
        .collect::<Result<Vec<_>, TransmissionError>>()?;

    // Sort by peak_snr descending
    results.sort_by(|(_, a), (_, b)| {
        b.peak_snr
            .partial_cmp(&a.peak_snr)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    Ok(results)
}

#[cfg(test)]
mod tests {
    use super::*;
    use nereids_endf::resonance::test_support::synthetic_isotope;

    // --- Offline synthetic tests (no network required) ---

    /// Peak selection and SNR computation with synthetic data (offline).
    #[test]
    fn test_synthetic_peak_snr() {
        // Matrix: weak resonance at 20 eV
        let matrix = synthetic_isotope(74, 182, 20.0, 1e-3, 0.05);
        // Trace: strong resonance at 8 eV (should dominate the |ΔT| spectrum)
        let trace = synthetic_isotope(72, 178, 8.0, 0.5, 0.05);

        let energies: Vec<f64> = (0..500).map(|i| 1.0 + (i as f64) * 49.0 / 499.0).collect();

        let config = TraceDetectabilityConfig {
            matrix_isotopes: &[(matrix.clone(), 6e-3)],
            energies: &energies,
            i0: 1000.0,
            temperature_k: 293.6,
            resolution: None,
            snr_threshold: 3.0,
        };

        let report = trace_detectability(&config, &trace, 1000.0).unwrap();

        // Basic sanity: spectrum has correct length
        assert_eq!(report.delta_t_spectrum.len(), energies.len());
        assert_eq!(report.energies.len(), energies.len());

        // Peak should be near the trace resonance at 8 eV (not the matrix at 20 eV)
        assert!(
            report.peak_energy_ev > 5.0 && report.peak_energy_ev < 12.0,
            "Peak energy should be near 8 eV, got {:.1} eV",
            report.peak_energy_ev,
        );

        // With a strong trace resonance at 1000 ppm, SNR should be well above 3
        assert!(
            report.peak_snr > 3.0,
            "Expected detectable (SNR > 3), got SNR={:.2}",
            report.peak_snr,
        );
        assert!(report.detectable);

        // peak_delta_t_per_ppm should be positive and consistent
        assert!(report.peak_delta_t_per_ppm > 0.0);
    }

    /// Survey sorting with synthetic isotopes (offline).
    #[test]
    fn test_synthetic_survey_sorting() {
        let matrix = synthetic_isotope(74, 182, 20.0, 1e-3, 0.05);
        // Strong trace (big resonance at 8 eV)
        let strong_trace = synthetic_isotope(72, 178, 8.0, 0.5, 0.05);
        // Weak trace (tiny resonance at 30 eV)
        let weak_trace = synthetic_isotope(26, 56, 30.0, 1e-5, 0.01);

        let energies: Vec<f64> = (0..500).map(|i| 1.0 + (i as f64) * 49.0 / 499.0).collect();

        let config = TraceDetectabilityConfig {
            matrix_isotopes: &[(matrix.clone(), 6e-3)],
            energies: &energies,
            i0: 1000.0,
            temperature_k: 293.6,
            resolution: None,
            snr_threshold: 3.0,
        };

        // Pass weak first — survey should reorder by SNR descending
        let results =
            trace_detectability_survey(&config, &[weak_trace, strong_trace], 500.0).unwrap();

        assert_eq!(results.len(), 2);
        assert!(
            results[0].1.peak_snr >= results[1].1.peak_snr,
            "Survey should sort by SNR descending: {:.2} >= {:.2}",
            results[0].1.peak_snr,
            results[1].1.peak_snr,
        );
        // The strong isotope (Z=72) should rank first
        assert!(
            results[0].0.starts_with("Hf"),
            "Strong trace (Hf) should rank first, got '{}'",
            results[0].0,
        );
    }

    /// Multi-matrix baseline: two matrix isotopes produce a deeper baseline
    /// than either alone, reducing trace detectability.
    #[test]
    fn test_multi_matrix_baseline() {
        // Two matrix isotopes with resonances at different energies
        let matrix_a = synthetic_isotope(74, 182, 20.0, 1e-3, 0.05);
        let matrix_b = synthetic_isotope(26, 56, 35.0, 2e-3, 0.04);
        let trace = synthetic_isotope(72, 178, 8.0, 0.5, 0.05);

        let energies: Vec<f64> = (0..500).map(|i| 1.0 + (i as f64) * 49.0 / 499.0).collect();

        // Single-matrix config
        let config_single = TraceDetectabilityConfig {
            matrix_isotopes: &[(matrix_a.clone(), 6e-3)],
            energies: &energies,
            i0: 1000.0,
            temperature_k: 293.6,
            resolution: None,
            snr_threshold: 3.0,
        };

        // Multi-matrix config (same total density split across two isotopes)
        let config_multi = TraceDetectabilityConfig {
            matrix_isotopes: &[(matrix_a, 3e-3), (matrix_b, 3e-3)],
            energies: &energies,
            i0: 1000.0,
            temperature_k: 293.6,
            resolution: None,
            snr_threshold: 3.0,
        };

        let report_single = trace_detectability(&config_single, &trace, 1000.0).unwrap();
        let report_multi = trace_detectability(&config_multi, &trace, 1000.0).unwrap();

        // Both should produce valid spectra of the correct length
        assert_eq!(report_single.delta_t_spectrum.len(), energies.len());
        assert_eq!(report_multi.delta_t_spectrum.len(), energies.len());

        // Multi-matrix should still detect the trace (strong resonance at 8 eV)
        assert!(
            report_multi.peak_snr > 0.0,
            "Multi-matrix SNR should be positive"
        );

        // The two configs produce different SNR values (different baselines)
        assert!(
            (report_single.peak_snr - report_multi.peak_snr).abs() > 1e-10,
            "Single vs multi-matrix should produce different SNR values"
        );
    }

    /// Opaque matrix: extremely high density causes exp(-n·σ) underflow.
    ///
    /// Reproduces #278: 1.0 at/barn matrix with a strong resonance gives
    /// T ≈ 0 (underflow), so ΔT = 0 and SNR = 0. The report should flag
    /// the matrix as opaque via `opaque_fraction > 0.5`.
    #[test]
    fn test_opaque_matrix_flags_high_opaque_fraction() {
        // Matrix with a strong resonance at 20 eV and very high density (1.0 at/barn)
        let matrix = synthetic_isotope(74, 182, 20.0, 0.5, 0.05);
        let trace = synthetic_isotope(72, 178, 8.0, 0.5, 0.05);

        let energies: Vec<f64> = (0..500).map(|i| 1.0 + (i as f64) * 49.0 / 499.0).collect();

        let config = TraceDetectabilityConfig {
            matrix_isotopes: &[(matrix, 1.0)], // 1.0 at/barn — extremely thick
            energies: &energies,
            i0: 10000.0,
            temperature_k: 293.6,
            resolution: None,
            snr_threshold: 3.0,
        };

        let report = trace_detectability(&config, &trace, 2350.0).unwrap();

        // With such a thick matrix, most bins should be opaque
        assert!(
            report.opaque_fraction > 0.5,
            "Expected opaque_fraction > 0.5 for 1.0 at/barn matrix, got {:.3}",
            report.opaque_fraction,
        );

        // SNR should be 0 or very close to 0 due to underflow
        assert!(
            report.peak_snr < 1.0,
            "Expected near-zero SNR for opaque matrix, got {:.2}",
            report.peak_snr,
        );
    }

    /// Normal density: opaque_fraction should be 0 or very small.
    #[test]
    fn test_normal_density_not_opaque() {
        let matrix = synthetic_isotope(74, 182, 20.0, 1e-3, 0.05);
        let trace = synthetic_isotope(72, 178, 8.0, 0.5, 0.05);

        let energies: Vec<f64> = (0..500).map(|i| 1.0 + (i as f64) * 49.0 / 499.0).collect();

        let config = TraceDetectabilityConfig {
            matrix_isotopes: &[(matrix, 6e-3)], // typical thin sample
            energies: &energies,
            i0: 1000.0,
            temperature_k: 293.6,
            resolution: None,
            snr_threshold: 3.0,
        };

        let report = trace_detectability(&config, &trace, 1000.0).unwrap();

        assert!(
            report.opaque_fraction < 0.01,
            "Expected opaque_fraction < 0.01 for normal density, got {:.3}",
            report.opaque_fraction,
        );
    }
}
