//! Element-count bounds, DBE check, and isotopic-pattern inference.
//!
//! This module also contains the `deduce_isotopic_pattern_inner` helper moved from
//! `algorithms.rs`, which infers plausible counts of C / S / Cl / Br from an MS1
//! isotope pattern.

#[cfg(feature = "serde")]
use serde::Deserialize;
use std::collections::HashMap;

use crate::element::{ELEMENT_ISOTOPES, NUM_ELEMENTS};
use crate::error::DecompositionError;
use crate::formula::Formula;
use crate::scoring::MassTolerance;

const MASS_ACCURACY_PPM_TO_DA_THRESHOLD: f64 = 200.0;

pub const MIN_MASS_FOR_TOLERANCE: f64 = 200.0;

/// Compute the double-bond equivalent (DBE) of a formula and check it against
/// the supplied bounds.
///
/// The DBE formula is:
///   DBE = C + (N + P) / 2 - (H + F + Cl + Br + I + Na + K) / 2 + 1
pub fn check_dbe(formula: &Formula, min_dbe: f64, max_dbe: f64, allow_half_integer: bool) -> bool {
    let c_count = formula[1];
    let n_count = formula[2];
    let p_count = formula[6];
    let h_count = formula[0];
    let na_count = formula[5];
    let k_count = formula[9];
    let f_count = formula[4];
    let cl_count = formula[8];
    let br_count = formula[10];
    let i_count = formula[11];

    let dbe = ((c_count) as f64) + ((n_count + p_count) as f64 / 2.0)
        - ((h_count + f_count + cl_count + br_count + i_count + na_count + k_count) as f64 / 2.0)
        + 1.0;

    if dbe < min_dbe || dbe > max_dbe {
        return false;
    }

    if allow_half_integer {
        (dbe * 2.0).fract() == 0.0
    } else {
        dbe.fract() == 0.0
    }
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Deserialize))]
pub struct DecompositionParams {
    pub tolerance_ppm: f64,
    pub min_dbe: f64,
    pub max_dbe: f64,
    pub allow_half_integer: bool,
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Deserialize))]
pub struct DecompositionKwargs {
    /// Optional relative tolerance in parts-per-million. Mutually exclusive
    /// with `tolerance_da`; exactly one of the two must be `Some`.
    pub tolerance_ppm: Option<f64>,
    /// Optional absolute tolerance in Daltons. Mutually exclusive with
    /// `tolerance_ppm`; exactly one of the two must be `Some`.
    pub tolerance_da: Option<f64>,
    pub min_dbe: f64,
    pub max_dbe: f64,
    pub dbe_mode: String,
    pub min_bounds: Option<Formula>,
    pub max_bounds: Option<Formula>,
}

impl DecompositionKwargs {
    /// Resolve the tolerance pair into a single [`MassTolerance`].
    ///
    /// Returns `DecompositionError::InvalidInput` if neither or both of
    /// `tolerance_ppm` and `tolerance_da` are set.
    pub fn to_mass_tolerance(&self) -> Result<MassTolerance, DecompositionError> {
        match (self.tolerance_ppm, self.tolerance_da) {
            (Some(ppm), None) => Ok(MassTolerance::Ppm(ppm)),
            (None, Some(da)) => Ok(MassTolerance::Dalton(da)),
            (Some(ppm), Some(da)) => Err(DecompositionError::InvalidInput(format!(
                "exactly one of tolerance_ppm ({ppm}) and tolerance_da ({da}) must be set, not both"
            ))),
            (None, None) => Err(DecompositionError::InvalidInput(
                "exactly one of tolerance_ppm and tolerance_da must be set, got neither"
                    .to_string(),
            )),
        }
    }
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Deserialize))]
pub struct SpectrumDecompositionParams {
    /// Mass tolerance for the spectrum decomposition. Choose either Daltons
    /// or ppm; the type system prevents passing both. When `Dalton` is
    /// supplied, it is converted to ppm per mass using the 200 Da rule
    /// (`max(mass, 200 Da)`).
    pub mass_tolerance: MassTolerance,
    pub min_dbe: f64,
    pub max_dbe: f64,
    pub allow_half_integer: bool,
    pub water_absorption: bool,
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Deserialize))]
pub struct CleanAndNormalizeSpectrumKwargs {
    /// Optional relative tolerance in ppm for the raw-fragment decomposition.
    /// Mutually exclusive with `raw_fragment_tolerance_da`; exactly one of
    /// the two must be `Some`.
    pub raw_fragment_tolerance_ppm: Option<f64>,
    /// Optional absolute tolerance in Daltons for the raw-fragment
    /// decomposition. Mutually exclusive with `raw_fragment_tolerance_ppm`.
    pub raw_fragment_tolerance_da: Option<f64>,
    /// Optional relative tolerance in ppm for the normalized-fragment
    /// decomposition. Mutually exclusive with
    /// `normalized_fragment_tolerance_da`; exactly one of the two must be
    /// `Some`.
    pub normalized_fragment_tolerance_ppm: Option<f64>,
    /// Optional absolute tolerance in Daltons for the normalized-fragment
    /// decomposition. Mutually exclusive with
    /// `normalized_fragment_tolerance_ppm`.
    pub normalized_fragment_tolerance_da: Option<f64>,
    pub min_dbe: f64,
    pub max_dbe: f64,
    pub dbe_mode: String,
    pub water_absorption: bool,
}

impl CleanAndNormalizeSpectrumKwargs {
    /// Resolve the raw-fragment tolerance pair into a [`MassTolerance`].
    ///
    /// Returns `DecompositionError::InvalidInput` if neither or both of
    /// `raw_fragment_tolerance_ppm` and `raw_fragment_tolerance_da` are set.
    pub fn raw_mass_tolerance(&self) -> Result<MassTolerance, DecompositionError> {
        match (self.raw_fragment_tolerance_ppm, self.raw_fragment_tolerance_da) {
            (Some(ppm), None) => Ok(MassTolerance::Ppm(ppm)),
            (None, Some(da)) => Ok(MassTolerance::Dalton(da)),
            (Some(ppm), Some(da)) => Err(DecompositionError::InvalidInput(format!(
                "exactly one of raw_fragment_tolerance_ppm ({ppm}) and raw_fragment_tolerance_da ({da}) must be set, not both"
            ))),
            (None, None) => Err(DecompositionError::InvalidInput(
                "exactly one of raw_fragment_tolerance_ppm and raw_fragment_tolerance_da must be set, got neither"
                    .to_string(),
            )),
        }
    }

    /// Resolve the normalized-fragment tolerance pair into a [`MassTolerance`].
    ///
    /// Returns `DecompositionError::InvalidInput` if neither or both of
    /// `normalized_fragment_tolerance_ppm` and
    /// `normalized_fragment_tolerance_da` are set.
    pub fn normalized_mass_tolerance(&self) -> Result<MassTolerance, DecompositionError> {
        match (
            self.normalized_fragment_tolerance_ppm,
            self.normalized_fragment_tolerance_da,
        ) {
            (Some(ppm), None) => Ok(MassTolerance::Ppm(ppm)),
            (None, Some(da)) => Ok(MassTolerance::Dalton(da)),
            (Some(ppm), Some(da)) => Err(DecompositionError::InvalidInput(format!(
                "exactly one of normalized_fragment_tolerance_ppm ({ppm}) and normalized_fragment_tolerance_da ({da}) must be set, not both"
            ))),
            (None, None) => Err(DecompositionError::InvalidInput(
                "exactly one of normalized_fragment_tolerance_ppm and normalized_fragment_tolerance_da must be set, got neither"
                    .to_string(),
            )),
        }
    }
}

#[derive(Debug, Clone)]
pub struct CorrectedFragment {
    pub normalized_mass: f64,
    pub intensity: f64,
    pub formula: Formula,
    pub error_ppm: f64,
}

pub struct CleanedAndNormalizedSpectrumResult {
    pub fragments: Vec<CorrectedFragment>,
}

#[derive(Debug, Clone)]
pub struct DecomposedFragment {
    pub formula: Formula,
    pub mass: f64,
    pub error_ppm: f64,
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Deserialize))]
pub struct CalibratedIsotopicModel {
    pub alpha: f64,
    pub beta: f64,
    pub offset: f64,
    pub hetero_lower_a: f64,
    pub hetero_lower_b: f64,
    pub hetero_upper_a: f64,
    pub hetero_upper_b: f64,
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Deserialize))]
pub struct IsotopicPatternParams {
    pub ms1_mass_tolerance_ppm: f64,
    pub isotopic_mass_tolerance_ppm: f64,
    pub minimum_intensity: f64,
    pub model: CalibratedIsotopicModel,
}

#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Deserialize))]
pub struct DeduceIsotopicPatternKwargs {
    pub ms1_mass_tolerance_ppm: f64,
    pub isotopic_mass_tolerance_ppm: f64,
    pub minimum_intensity: f64,
    pub model: CalibratedIsotopicModel,
    pub min_bounds: Option<HashMap<String, i32>>,
    pub max_bounds: Option<HashMap<String, i32>>,
}

pub fn default_min_bounds() -> Formula {
    [0; NUM_ELEMENTS]
}

pub fn default_max_bounds() -> Formula {
    let mut bounds = [0; NUM_ELEMENTS];
    // H=0, C=1, N=2, O=3, F=4, Na=5, P=6, S=7, Cl=8, K=9, Br=10, I=11
    // Based on python DEFAULT_MAX_BOUND
    bounds[0] = 100; // H
    bounds[1] = 50; // C
    bounds[2] = 20; // N
    bounds[3] = 20; // O
    bounds[4] = 40; // F
    bounds[5] = 0; // Na
    bounds[6] = 5; // P
    bounds[7] = 5; // S
    bounds[8] = 10; // Cl
    bounds[9] = 0; // K
    bounds[10] = 10; // Br
    bounds[11] = 5; // I
    bounds
}

/// Infer C / S / Cl / Br count bounds from an MS1 isotope pattern.
///
/// Returns `[C_low, S_low, Cl_low, Br_low, C_high, S_high, Cl_high, Br_high]`.
/// If no precursor is found at `precursor_mz`, returns `[-1.0; 8]`.
pub fn deduce_isotopic_pattern_inner(
    precursor_mz: f64,
    ms1_mzs: &[f64],
    ms1_intensities: &[f64],
    params: &IsotopicPatternParams,
) -> [f64; 8] {
    let ms1_absolute_tolerance =
        precursor_mz.max(MASS_ACCURACY_PPM_TO_DA_THRESHOLD) * params.ms1_mass_tolerance_ppm * 1e-6;
    let isotopic_absolute_tolerance = precursor_mz.max(MASS_ACCURACY_PPM_TO_DA_THRESHOLD)
        * params.isotopic_mass_tolerance_ppm
        * 1e-6;

    // Find precursor
    let precursor_indices: Vec<usize> = ms1_mzs
        .iter()
        .enumerate()
        .filter(|&(_, &mz)| (mz - precursor_mz).abs() <= ms1_absolute_tolerance)
        .map(|(i, _)| i)
        .collect();

    if precursor_indices.is_empty() {
        return [-1.0; 8];
    }

    // Find precursor intensity (max intensity among matches)
    let mut precursor_ms1_intensity = 0.0;
    let mut best_precursor_idx = 0;
    for &idx in &precursor_indices {
        if ms1_intensities[idx] > precursor_ms1_intensity {
            precursor_ms1_intensity = ms1_intensities[idx];
            best_precursor_idx = idx;
        }
    }
    let precursor_ms1_mz = ms1_mzs[best_precursor_idx];

    let mut result = [0.0; 8];

    // Element indices in ELEMENT_SYMBOLS
    // C: 1, S: 7, Cl: 8, Br: 10
    // Output indices: 0: C_low, 1: S_low, 2: Cl_low, 3: Br_low, 4: C_high, 5: S_high, 6: Cl_high, 7: Br_high

    let target_elements = [
        (1, 0),  // (Element Index, Result Index) -> C
        (7, 1),  // S
        (8, 2),  // Cl
        (10, 3), // Br
    ];

    for &(elem_idx, res_idx) in &target_elements {
        let iso_props = ELEMENT_ISOTOPES[elem_idx].expect("Element should have isotopic props");

        let peak_mz = precursor_ms1_mz + iso_props.mass_diff;
        let peak_indices: Vec<usize> = ms1_mzs
            .iter()
            .enumerate()
            .filter(|&(_, &mz)| (mz - peak_mz).abs() <= isotopic_absolute_tolerance)
            .map(|(i, _)| i)
            .collect();

        let peak_total_intensity = if peak_indices.is_empty() {
            0.0
        } else {
            peak_indices
                .iter()
                .map(|&i| ms1_intensities[i])
                .fold(0.0, f64::max)
        };

        let lower;
        let upper;

        if peak_total_intensity < params.minimum_intensity {
            lower = 0.0;
            upper = (params.minimum_intensity * iso_props.prob_0)
                / (iso_props.prob_1 * precursor_ms1_intensity);
        } else {
            let observed = peak_total_intensity;
            let expected_hat =
                params.model.alpha * (observed + params.model.offset).powf(params.model.beta);
            let w_l = params.model.hetero_lower_a
                * (observed + params.model.offset).powf(params.model.hetero_lower_b);
            let w_u = params.model.hetero_upper_a
                * (observed + params.model.offset).powf(params.model.hetero_upper_b);
            let expected_lower = expected_hat * (-w_l).exp();
            let expected_upper = expected_hat * w_u.exp();
            lower =
                (expected_lower * iso_props.prob_0) / (iso_props.prob_1 * precursor_ms1_intensity);
            upper =
                (expected_upper * iso_props.prob_0) / (iso_props.prob_1 * precursor_ms1_intensity);
        }

        result[res_idx] = lower;
        result[res_idx + 4] = upper;
    }

    result
}

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

    #[test]
    fn test_check_dbe() {
        // C6H12O6 (glucose) has DBE = 6 - 12/2 + 0/2 + 1 = 1
        let glucose = [12, 6, 0, 6, 0, 0, 0, 0, 0, 0, 0, 0];
        let dbe = 6.0 - 6.0 + 0.0 + 1.0;
        assert!(check_dbe(&glucose, 0.0, dbe + 0.5, false));
        assert!(check_dbe(&glucose, dbe, dbe, false));
        assert!(!check_dbe(&glucose, dbe + 0.1, 100.0, false));
    }
}
