//! Public API: types, configuration, candidate ranking, and entry-point functions.
//!
//! This module wraps the lower-level `MassDecomposer` and exposes a small,
//! easy-to-use surface area for callers (e.g. `spectral-autoencoder`).
//!
//! Determinism: candidates are always returned sorted by score descending, with
//! ties broken by lexicographic formula string, then truncated to
//! `config.candidate_limit`. For identical inputs the output is byte-for-byte
//! identical.

use rayon::prelude::*;
use std::collections::HashMap;
use std::sync::Arc;

use crate::bounds::{
    DecompositionParams, IsotopicPatternParams, MIN_MASS_FOR_TOLERANCE,
    deduce_isotopic_pattern_inner, default_max_bounds, default_min_bounds,
};
use crate::decomposer::MassDecomposer;
use crate::element::{ELEMENT_SYMBOLS, NUM_ELEMENTS};
use crate::error::DecompositionError;
use crate::formula::{
    Formula, formula_dbe, formula_mass, formula_mass_defect, formula_nominal_mass,
    formula_to_composition, formula_to_string, is_valid_valence,
};

// Re-export the public formula types so callers can use them from this module.
pub use crate::formula::{Composition, Element};

/// Mass tolerance for a decomposition run.
///
/// The tolerance can be specified either as an absolute error in Daltons or as
/// a relative error in parts-per-million. Exactly one variant must be chosen;
/// the type system prevents passing both at the same time.
#[derive(Debug, Clone, Copy, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize))]
pub enum MassTolerance {
    /// Absolute tolerance in Daltons.
    Dalton(f64),
    /// Relative tolerance in parts-per-million.
    Ppm(f64),
}

impl MassTolerance {
    /// Convert this tolerance to a parts-per-million value for the given
    /// target mass, applying the 200 Da rule: the reference mass used for the
    /// conversion is `max(target_mass, 200 Da)`.
    pub fn to_ppm(self, target_mass: f64) -> f64 {
        match self {
            MassTolerance::Dalton(da) => da / target_mass.max(MIN_MASS_FOR_TOLERANCE) * 1e6,
            MassTolerance::Ppm(ppm) => ppm,
        }
    }
}

/// Isotopic pattern supplied by the caller for the precursor.
#[derive(Debug, Clone)]
pub struct IsotopicPattern {
    /// (m/z, intensity) peaks.
    pub peaks: Vec<(f64, f64)>,
}

/// Precursor information.
#[derive(Debug, Clone)]
pub struct PrecursorInfo {
    pub monoisotopic_mass: f64,
    pub charge: i32,
    pub adduct: Option<String>,
    pub isotopic_pattern: Option<IsotopicPattern>,
}

/// A single fragment peak.
#[derive(Debug, Clone, Copy)]
pub struct FragmentPeak {
    pub mz: f64,
    pub intensity: f64,
}

/// A decomposed molecular formula.
#[derive(Debug, Clone)]
pub struct MolecularFormula {
    pub composition: Composition,
    pub monoisotopic_mass: f64,
    pub nominal_mass: u32,
    pub dbe: f64,
    pub mass_defect: f64,
    pub is_valid_valence: bool,
}

impl MolecularFormula {
    pub fn from_internal(f: &Formula) -> Self {
        let composition = formula_to_composition(f);
        let monoisotopic_mass = formula_mass(f);
        let nominal_mass = formula_nominal_mass(f);
        let dbe = formula_dbe(f);
        let mass_defect = formula_mass_defect(f);
        let is_valid_valence = is_valid_valence(f);
        Self {
            composition,
            monoisotopic_mass,
            nominal_mass,
            dbe,
            mass_defect,
            is_valid_valence,
        }
    }

    /// Convert this formula into a fixed-width feature vector.
    ///
    /// The output has length `config.feature_count` and includes the features
    /// specified in `config` in a fixed, deterministic order.
    pub fn to_feature_vector(&self, intensity: f32, config: &FormulaFeatureConfig) -> Vec<f32> {
        let mut out = Vec::with_capacity(config.feature_count);
        out.push(self.monoisotopic_mass as f32);
        out.push(intensity);
        out.push(self.dbe as f32);
        out.push(self.mass_defect as f32);
        for elem in &config.include_element_counts {
            let count = self.composition.counts.get(elem).copied().unwrap_or(0) as f32;
            out.push(count);
        }
        for (a, b) in &config.include_ratios {
            let ca = self.composition.counts.get(a).copied().unwrap_or(0) as f32;
            let cb = self.composition.counts.get(b).copied().unwrap_or(0) as f32;
            let ratio = if cb > 0.0 { ca / cb } else { 0.0 };
            out.push(ratio);
        }
        debug_assert_eq!(
            out.len(),
            config.feature_count,
            "feature vector length mismatch: expected {}, got {}",
            config.feature_count,
            out.len()
        );
        out
    }
}

/// Configuration for one decomposition run.
#[derive(Debug, Clone)]
pub struct DecompositionConfig {
    /// Mass tolerance for candidate filtering. Choose either an absolute
    /// tolerance in Daltons (`MassTolerance::Dalton`) or a relative tolerance
    /// in parts-per-million (`MassTolerance::Ppm`). The two units are not
    /// interchangeable; the type system forces callers to pick exactly one.
    ///
    /// Internally, a Dalton tolerance is converted to ppm using the 200 Da
    /// rule: the reference mass is `max(target_mass, 200 Da)`. This avoids
    /// excessively tight ppm tolerances at very low masses.
    pub mass_tolerance: MassTolerance,
    /// Allowed elements. Anything not in this list is constrained to 0.
    pub allowed_elements: Vec<Element>,
    /// Optional per-element maximum counts. If `None`, uses
    /// `default_max_bounds()` from the `bounds` module.
    pub max_element_counts: Option<Composition>,
    /// DBE lower bound (inclusive).
    pub min_dbe: f64,
    /// DBE upper bound (inclusive).
    pub max_dbe: f64,
    /// Maximum number of candidates returned per mass.
    pub candidate_limit: usize,
    /// Allow half-integer DBE values.
    pub allow_half_integer: bool,
    /// Isotopic-pattern parameters used for isotope-aware scoring.
    ///
    /// If `None`, isotope-aware scoring is disabled and only mass is used.
    pub isotopic_pattern_params: Option<IsotopicPatternParams>,
}

impl DecompositionConfig {
    /// Sensible default for positive-ESI LC-MS/MS of small molecules
    /// (metabolomics-like workloads). Elements: H, C, N, O, P, S.
    pub fn default_lcms_positive() -> Self {
        Self {
            mass_tolerance: MassTolerance::Dalton(0.01),
            allowed_elements: vec![
                "H".to_string(),
                "C".to_string(),
                "N".to_string(),
                "O".to_string(),
                "P".to_string(),
                "S".to_string(),
            ],
            max_element_counts: None,
            min_dbe: 0.0,
            max_dbe: 40.0,
            candidate_limit: 1,
            allow_half_integer: false,
            isotopic_pattern_params: None,
        }
    }
}

/// Result for one fragment peak.
#[derive(Debug, Clone)]
pub struct FragmentFormulaResult {
    pub peak: FragmentPeak,
    pub candidates: Vec<(MolecularFormula, f64)>,
}

impl FragmentFormulaResult {
    /// Convenience accessor for the single-best formula.
    pub fn best_formula(&self) -> Option<&MolecularFormula> {
        self.candidates.first().map(|(f, _)| f)
    }
}

/// Configuration for the formula feature vector.
#[derive(Debug, Clone)]
pub struct FormulaFeatureConfig {
    pub feature_count: usize,
    pub include_element_counts: Vec<Element>,
    pub include_ratios: Vec<(Element, Element)>,
}

impl FormulaFeatureConfig {
    /// Default 12-feature vector: mass, intensity, DBE, mass defect,
    /// C, H, N, O, P, S counts, H/C and N/O ratios.
    pub fn default_12() -> Self {
        let include_element_counts = vec![
            "C".to_string(),
            "H".to_string(),
            "N".to_string(),
            "O".to_string(),
            "P".to_string(),
            "S".to_string(),
        ];
        let include_ratios = vec![
            ("H".to_string(), "C".to_string()),
            ("N".to_string(), "O".to_string()),
        ];
        // 4 leading + 6 counts + 2 ratios = 12
        Self {
            feature_count: 12,
            include_element_counts,
            include_ratios,
        }
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Convert `allowed_elements` and optional `max_element_counts` to internal
/// `Formula` min/max bounds. Elements not in `allowed_elements` are forced to 0.
fn build_bounds(
    allowed_elements: &[Element],
    max_element_counts: Option<&Composition>,
) -> (Formula, Formula) {
    let allowed_set: std::collections::HashSet<&str> =
        allowed_elements.iter().map(|s| s.as_str()).collect();
    let default_max = default_max_bounds();
    let mut min_b = default_min_bounds();
    let mut max_b = default_max;
    for i in 0..NUM_ELEMENTS {
        if !allowed_set.contains(ELEMENT_SYMBOLS[i]) {
            min_b[i] = 0;
            max_b[i] = 0;
        } else if let Some(max_counts) = max_element_counts {
            // Override default max with user-supplied value.
            max_b[i] = max_counts
                .counts
                .get(ELEMENT_SYMBOLS[i])
                .copied()
                .unwrap_or(0) as i32;
        }
    }
    (min_b, max_b)
}

/// Compute the predicted M+1 / M+0 intensity ratio for a formula.
///
/// This is a coarse approximation using only the most abundant heavy isotopes:
/// C-13, N-15, S-33, Cl-37, Br-81. It is good enough for ranking.
fn predicted_m1_ratio(f: &Formula) -> f64 {
    let c = f[1] as f64 * 0.0107;
    let n = f[2] as f64 * 0.004;
    let s = f[7] as f64 * 0.0429;
    let cl = f[8] as f64 * 0.2422;
    let br = f[10] as f64 * 0.4931;
    c + n + s + cl + br
}

/// Score a candidate against the supplied pattern (if any) and return the
/// adjusted score. A higher score is better.
fn score_with_isotope(formula: &Formula, error_ppm: f64, pattern: Option<&IsotopicPattern>) -> f64 {
    let base = -error_ppm.abs();
    let Some(pattern) = pattern else {
        return base;
    };
    // Use the first two peaks (M+0 and M+1) for a coarse check.
    if pattern.peaks.len() < 2 {
        return base;
    }
    let m0 = pattern.peaks[0].1;
    let m1 = pattern.peaks[1].1;
    if m0 <= 0.0 {
        return base;
    }
    let observed = m1 / m0;
    let predicted = predicted_m1_ratio(formula);
    // Add a small bonus for isotope agreement, weighted at 0.5 so it never
    // overrides a 1+ ppm mass error.
    let iso_bonus = -0.5 * (observed - predicted).abs();
    base + iso_bonus
}

/// Rank and truncate a list of candidates. Ordering is:
/// 1. score descending
/// 2. formula string ascending (lexicographic tie-break)
/// 3. original index ascending (final deterministic tie-break)
fn rank_candidates(
    candidates: Vec<(Formula, f64, f64)>, // (formula, error_ppm, score)
    limit: usize,
) -> Vec<(Formula, f64, f64)> {
    let mut indexed: Vec<(usize, (Formula, f64, f64))> =
        candidates.into_iter().enumerate().collect();
    indexed.sort_by(|a, b| {
        let sa = a.1.2;
        let sb = b.1.2;
        // score descending
        sb.partial_cmp(&sa)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| {
                let fa = formula_to_string(&a.1.0);
                let fb = formula_to_string(&b.1.0);
                fa.cmp(&fb)
            })
            .then_with(|| a.0.cmp(&b.0))
    });
    indexed.truncate(limit);
    indexed.into_iter().map(|(_, x)| x).collect()
}

/// Decompose the precursor mass into candidate formulas.
///
/// Returns up to `config.candidate_limit` candidates sorted by score descending.
/// If `precursor.isotopic_pattern` is set and `config.isotopic_pattern_params` is
/// set, candidates that violate the inferred C / S / Cl / Br bounds are filtered
/// out and a small isotope-match bonus is added to the score.
pub fn decompose_precursor(
    precursor: &PrecursorInfo,
    config: &DecompositionConfig,
) -> Result<Vec<(MolecularFormula, f64)>, DecompositionError> {
    if precursor.monoisotopic_mass <= 0.0 || !precursor.monoisotopic_mass.is_finite() {
        return Err(DecompositionError::InvalidInput(format!(
            "monoisotopic_mass must be positive and finite, got {}",
            precursor.monoisotopic_mass
        )));
    }
    if precursor.charge != 1 && precursor.charge != -1 {
        return Err(DecompositionError::InvalidInput(format!(
            "only charge +1 or -1 is supported, got {}",
            precursor.charge
        )));
    }

    let (min_b, max_b) = build_bounds(&config.allowed_elements, config.max_element_counts.as_ref());
    let decomposer = MassDecomposer::new(min_b, max_b)?;
    let params = DecompositionParams {
        tolerance_ppm: config.mass_tolerance.to_ppm(precursor.monoisotopic_mass),
        min_dbe: config.min_dbe,
        max_dbe: config.max_dbe,
        allow_half_integer: config.allow_half_integer,
    };

    let (formulas, errors_ppm) = decomposer.decompose(precursor.monoisotopic_mass, &params);

    if formulas.is_empty() {
        return Err(DecompositionError::NoCandidate);
    }

    // Isotope-aware filtering and scoring.
    let iso_bounds = if let (Some(pattern), Some(iso_params)) = (
        precursor.isotopic_pattern.as_ref(),
        config.isotopic_pattern_params.as_ref(),
    ) {
        let mzs: Vec<f64> = pattern.peaks.iter().map(|p| p.0).collect();
        let intensities: Vec<f64> = pattern.peaks.iter().map(|p| p.1).collect();
        let deduced = deduce_isotopic_pattern_inner(
            precursor.monoisotopic_mass,
            &mzs,
            &intensities,
            iso_params,
        );
        Some(deduced)
    } else {
        None
    };

    // Apply isotope filtering, then compute the final score per candidate.
    let mut candidates: Vec<(Formula, f64, f64)> = Vec::with_capacity(formulas.len());
    for (f, &error_ppm) in formulas.iter().zip(errors_ppm.iter()) {
        if let Some(deduced) = iso_bounds {
            // Element indices for C, S, Cl, Br are 1, 7, 8, 10.
            // deduced: [C_l, S_l, Cl_l, Br_l, C_u, S_u, Cl_u, Br_u]
            // The values are continuous floats; we apply the same rounding
            // the legacy Polars interface uses: ceil for lower, floor for upper.
            let c_l = deduced[0].ceil() as i32;
            let s_l = deduced[1].ceil() as i32;
            let cl_l = deduced[2].ceil() as i32;
            let br_l = deduced[3].ceil() as i32;
            let c_u = deduced[4].floor() as i32;
            let s_u = deduced[5].floor() as i32;
            let cl_u = deduced[6].floor() as i32;
            let br_u = deduced[7].floor() as i32;
            if f[1] < c_l || f[1] > c_u {
                continue;
            }
            if f[7] < s_l || f[7] > s_u {
                continue;
            }
            if f[8] < cl_l || f[8] > cl_u {
                continue;
            }
            if f[10] < br_l || f[10] > br_u {
                continue;
            }
        }
        let score = score_with_isotope(f, error_ppm, precursor.isotopic_pattern.as_ref());
        candidates.push((*f, error_ppm, score));
    }

    if candidates.is_empty() {
        return Err(DecompositionError::NoCandidate);
    }

    let ranked = rank_candidates(candidates, config.candidate_limit);

    Ok(ranked
        .into_iter()
        .map(|(f, error_ppm, _score)| (MolecularFormula::from_internal(&f), error_ppm))
        .collect())
}

/// Internal helper used by both `decompose_fragments` and `decompose_fragment_batch`.
fn decompose_one_fragment(
    peak: FragmentPeak,
    config: &DecompositionConfig,
    decomposer: &MassDecomposer,
) -> FragmentFormulaResult {
    let params = DecompositionParams {
        tolerance_ppm: config.mass_tolerance.to_ppm(peak.mz),
        min_dbe: config.min_dbe,
        max_dbe: config.max_dbe,
        allow_half_integer: config.allow_half_integer,
    };
    let (formulas, errors_ppm) = decomposer.decompose(peak.mz, &params);
    let candidates: Vec<(Formula, f64, f64)> = formulas
        .into_iter()
        .zip(errors_ppm.into_iter())
        .map(|(f, e)| {
            let score = -e.abs();
            (f, e, score)
        })
        .collect();
    let ranked = rank_candidates(candidates, config.candidate_limit);
    let candidates = ranked
        .into_iter()
        .map(|(f, error_ppm, _score)| (MolecularFormula::from_internal(&f), error_ppm))
        .collect();
    FragmentFormulaResult { peak, candidates }
}

/// Decompose each fragment peak into candidate formulas, using the precursor
/// formula to constrain neutral losses.
///
/// Each fragment's per-element count cannot exceed the precursor's count for
/// that element (plus optional H2O allowance for water-absorption modes).
pub fn decompose_fragments(
    precursor: &MolecularFormula,
    fragments: &[FragmentPeak],
    config: &DecompositionConfig,
) -> Result<Vec<FragmentFormulaResult>, DecompositionError> {
    if fragments.is_empty() {
        return Ok(Vec::new());
    }
    let precursor_formula = composition_to_internal(&precursor.composition)?;
    let max_b = precursor_formula;
    let min_b = [0i32; NUM_ELEMENTS];
    let decomposer = MassDecomposer::new(min_b, max_b)?;

    let decomposer = Arc::new(decomposer);
    let mut out = Vec::with_capacity(fragments.len());
    for &peak in fragments {
        out.push(decompose_one_fragment(peak, config, &decomposer));
    }
    Ok(out)
}

/// Same as `decompose_fragments` but processes the fragment list in parallel
/// using rayon. The result order matches the input order.
pub fn decompose_fragment_batch(
    precursor: &MolecularFormula,
    fragments: &[FragmentPeak],
    config: &DecompositionConfig,
) -> Result<Vec<FragmentFormulaResult>, DecompositionError> {
    if fragments.is_empty() {
        return Ok(Vec::new());
    }
    let precursor_formula = composition_to_internal(&precursor.composition)?;
    let max_b = precursor_formula;
    let min_b = [0i32; NUM_ELEMENTS];
    let decomposer = Arc::new(MassDecomposer::new(min_b, max_b)?);

    let results: Vec<FragmentFormulaResult> = fragments
        .par_iter()
        .map(|&peak| decompose_one_fragment(peak, config, &decomposer))
        .collect();
    Ok(results)
}

fn composition_to_internal(comp: &Composition) -> Result<Formula, DecompositionError> {
    let mut f = [0i32; NUM_ELEMENTS];
    for (elem, count) in &comp.counts {
        let idx = ELEMENT_SYMBOLS
            .iter()
            .position(|&s| s == elem)
            .ok_or_else(|| {
                DecompositionError::InvalidInput(format!("unknown element symbol '{}'", elem))
            })?;
        f[idx] = *count as i32;
    }
    Ok(f)
}

/// Convenience helper: build a `MolecularFormula` from a `Composition`.
pub fn molecular_formula_from_composition(comp: Composition) -> MolecularFormula {
    let mut f = [0i32; NUM_ELEMENTS];
    for (elem, count) in &comp.counts {
        if let Some(idx) = ELEMENT_SYMBOLS.iter().position(|&s| s == elem) {
            f[idx] = *count as i32;
        }
    }
    MolecularFormula::from_internal(&f)
}

/// Build a `Composition` from explicit `(symbol, count)` pairs.
pub fn composition_from_pairs(pairs: &[(String, u32)]) -> Composition {
    let mut counts = HashMap::new();
    for (k, v) in pairs {
        counts.insert(k.clone(), *v);
    }
    Composition { counts }
}

// ---------------------------------------------------------------------------
// Re-exports for backward-compat with the existing crate
// ---------------------------------------------------------------------------

pub use crate::decomposer::SpectrumDecomposer as InternalSpectrumDecomposer;

#[allow(unused_imports)]
pub(crate) use crate::formula::formula_to_string as _formula_to_string_re_export;

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

    #[test]
    fn test_predicted_m1_ratio_zero() {
        let empty = [0i32; NUM_ELEMENTS];
        assert_eq!(predicted_m1_ratio(&empty), 0.0);
    }

    #[test]
    fn test_predicted_m1_ratio_ch4() {
        // CH4: C=1, H=4. Predicted: 0.0107 (only C contributes).
        let ch4 = [4, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let r = predicted_m1_ratio(&ch4);
        assert!((r - 0.0107).abs() < 1e-6);
    }

    #[test]
    fn test_rank_candidates_deterministic() {
        let cands = vec![
            ([0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 1.0, -1.0),
            ([2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 1.0, -1.0),
            ([0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 0.5, -0.5),
        ];
        let ranked = rank_candidates(cands, 10);
        // First should be the smallest error (CH2O with 0.5 ppm).
        assert_eq!(ranked[0].2, -0.5);
    }

    #[test]
    fn test_composition_from_pairs() {
        let c = composition_from_pairs(&[("C".into(), 6), ("H".into(), 12), ("O".into(), 6)]);
        assert_eq!(c.counts.get("C"), Some(&6));
        assert_eq!(c.counts.get("H"), Some(&12));
        assert_eq!(c.counts.get("O"), Some(&6));
    }

    #[test]
    fn test_default_lcms_positive() {
        let cfg = DecompositionConfig::default_lcms_positive();
        assert_eq!(cfg.allowed_elements, vec!["H", "C", "N", "O", "P", "S"]);
        assert_eq!(cfg.candidate_limit, 1);
        assert_eq!(cfg.mass_tolerance, MassTolerance::Dalton(0.01));
        assert!(!cfg.allow_half_integer);
    }

    #[test]
    fn test_feature_vector_default_12() {
        let cfg = FormulaFeatureConfig::default_12();
        assert_eq!(cfg.feature_count, 12);
        let formula = molecular_formula_from_composition(composition_from_pairs(&[
            ("C".into(), 6),
            ("H".into(), 12),
            ("O".into(), 6),
        ]));
        let v = formula.to_feature_vector(1.0, &cfg);
        assert_eq!(v.len(), 12);
    }

    #[test]
    fn test_invalid_precursor_mass() {
        let p = PrecursorInfo {
            monoisotopic_mass: -1.0,
            charge: 1,
            adduct: None,
            isotopic_pattern: None,
        };
        let cfg = DecompositionConfig::default_lcms_positive();
        assert!(decompose_precursor(&p, &cfg).is_err());
    }
}
