mzannotate 0.2.0

//! Handle fragment related issues, access provided if you want to dive deeply into fragments in
//! your own code.

use std::{
    cmp::Ordering,
    fmt::{Debug, Display},
    num::NonZeroU16,
    sync::LazyLock,
};

use itertools::Itertools;
use mzcore::{
    chemistry::{AmbiguousLabel, CachedCharge, ChargeRange, NeutralLoss, OutputMolecularFormula},
    molecular_formula,
    prelude::*,
    quantities::{Multi, Tolerance},
    system::{self, MassOverCharge, OrderedMassOverCharge, Ratio, isize::Charge},
};
use ordered_float::OrderedFloat;
use serde::{Deserialize, Serialize};
use thin_vec::ThinVec;

use crate::{annotation::model::PossiblePrimaryIons, fragment::FragmentType};

/// A theoretical fragment
#[derive(Debug, Deserialize, Hash, Serialize)]
pub struct Fragment<Mode: MassOutputMode> {
    /// The theoretical composition
    pub formula: Option<Mode::Output>,
    /// The charge
    pub charge: Charge,
    /// The annotation for this fragment
    pub ion: FragmentType,
    /// The isotope of this fragment (this needs to be added to the formula separately)
    pub isotope: ThinVec<(i32, Isotope)>,
    /// The peptidoform this fragment comes from, saved as the index into the list of peptidoform
    /// in the overarching [`mzcore::sequence::PeptidoformIonSet`] struct
    pub peptidoform_ion_index: Option<usize>,
    /// The peptide this fragment comes from, saved as the index into the list of peptides in the
    /// overarching [`mzcore::sequence::PeptidoformIon`] struct
    pub peptidoform_index: Option<usize>,
    /// Any neutral losses applied
    pub neutral_loss: ThinVec<NeutralLoss>,
    /// m/z deviation, if known (from mzPAF)
    pub deviation: Option<Tolerance<OrderedMassOverCharge>>,
    /// Confidence in this annotation (from mzPAF)
    pub confidence: Option<OrderedFloat<f64>>,
    /// If this is an auxiliary fragment (from mzPAF)
    pub auxiliary: bool,
}

impl<Mode: MassOutputMode> Clone for Fragment<Mode> {
    fn clone(&self) -> Self {
        Self {
            formula: self.formula.clone(),
            charge: self.charge.clone(),
            ion: self.ion.clone(),
            isotope: self.isotope.clone(),
            peptidoform_ion_index: self.peptidoform_ion_index.clone(),
            peptidoform_index: self.peptidoform_index.clone(),
            neutral_loss: self.neutral_loss.clone(),
            deviation: self.deviation.clone(),
            confidence: self.confidence.clone(),
            auxiliary: self.auxiliary.clone(),
        }
    }
}

impl<Mode: MassOutputMode> Default for Fragment<Mode> {
    fn default() -> Self {
        Self {
            formula: Default::default(),
            charge: Default::default(),
            ion: Default::default(),
            isotope: Default::default(),
            peptidoform_ion_index: Default::default(),
            peptidoform_index: Default::default(),
            neutral_loss: Default::default(),
            deviation: Default::default(),
            confidence: Default::default(),
            auxiliary: Default::default(),
        }
    }
}

impl<Mode: MassOutputMode> PartialOrd for Fragment<Mode> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<Mode: MassOutputMode> Ord for Fragment<Mode> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.formula
            .cmp(&other.formula)
            .then(self.charge.cmp(&other.charge))
            .then(self.ion.cmp(&other.ion))
            .then(self.isotope.cmp(&other.isotope))
            .then(self.peptidoform_ion_index.cmp(&other.peptidoform_ion_index))
            .then(self.peptidoform_index.cmp(&other.peptidoform_index))
            .then(self.neutral_loss.cmp(&other.neutral_loss))
            .then(self.deviation.cmp(&other.deviation))
            .then(self.confidence.cmp(&other.confidence))
            .then(self.auxiliary.cmp(&other.auxiliary))
    }
}

impl<Mode: MassOutputMode> PartialEq for Fragment<Mode> {
    fn eq(&self, other: &Self) -> bool {
        self.formula == other.formula
            && self.charge == other.charge
            && self.ion == other.ion
            && self.isotope == other.isotope
            && self.peptidoform_ion_index == other.peptidoform_ion_index
            && self.peptidoform_index == other.peptidoform_index
            && self.neutral_loss == other.neutral_loss
            && self.deviation == other.deviation
            && self.confidence == other.confidence
            && self.auxiliary == other.auxiliary
    }
}

impl<Mode: MassOutputMode> Eq for Fragment<Mode> {}

impl<Mode: MassOutputMode> Fragment<Mode> {
    /// Get the mz
    pub fn mz(&self, mode: MassMode) -> Option<MassOverCharge> {
        self.formula.as_ref().map(|f| {
            f.mass(mode) / system::f64::Charge::new::<system::charge::e>(self.charge.value as f64)
        })
    }

    /// Get the ppm difference between two fragments
    pub fn ppm(&self, other: &Self, mode: MassMode) -> Option<Ratio> {
        self.mz(mode)
            .and_then(|mz| other.mz(mode).map(|omz| (mz, omz)))
            .map(|(mz, omz)| mz.ppm(omz))
    }

    /// Create a new fragment
    #[must_use]
    pub fn new(
        theoretical_mass: Mode::Output,
        charge: Charge,
        peptidoform_ion_index: usize,
        peptidoform_index: usize,
        ion: FragmentType,
    ) -> Self {
        Self {
            formula: Some(theoretical_mass),
            charge,
            ion,
            isotope: ThinVec::new(),
            peptidoform_ion_index: Some(peptidoform_ion_index),
            peptidoform_index: Some(peptidoform_index),
            neutral_loss: ThinVec::new(),
            deviation: None,
            confidence: None,
            auxiliary: false,
        }
    }

    /// Generate a list of possible fragments from the list of possible preceding termini and
    /// neutral losses. Ignores any neutral loss that would result in a negative number of any
    /// element.
    /// # Panics
    /// When the charge range results in a negative charge
    #[expect(clippy::too_many_arguments)]
    #[must_use]
    pub fn generate_all(
        theoretical_mass: &Multi<Mode::Output>,
        peptidoform_ion_index: usize,
        peptidoform_index: usize,
        annotation: &FragmentType,
        termini: &Multi<Mode::Output>,
        neutral_losses: &[Vec<NeutralLoss>],
        charge_carriers: &mut CachedCharge,
        charge_range: ChargeRange,
    ) -> Vec<Self> {
        let charges = charge_carriers.range(charge_range);
        let losses = std::iter::once(None)
            .chain(neutral_losses.iter().map(Some))
            .collect::<Vec<_>>();
        let mut result = Vec::with_capacity(
            termini.len() * theoretical_mass.len() * charges.len() * losses.len(),
        );
        for term in termini.iter() {
            for mass in theoretical_mass.iter() {
                let f = term.clone() + mass.clone();
                for charge in &charges {
                    let f = f.clone()
                        + charge.calculate_mass_inner::<Mode>(
                            SequencePosition::default(),
                            peptidoform_index,
                        );
                    if f.contains_negative_amount() {
                        continue;
                    }
                    let z = Charge::new::<system::e>(charge.charge().value);
                    for loss in &losses {
                        let f = f.clone()
                            + Mode::from_formula(
                                loss.iter().flat_map(|l| l.iter()).sum::<MolecularFormula>(),
                            );
                        if f.contains_negative_amount() {
                            continue;
                        }
                        result.push(Self {
                            formula: Some(f),
                            charge: z,
                            ion: annotation.clone(),
                            isotope: ThinVec::new(),
                            peptidoform_ion_index: Some(peptidoform_ion_index),
                            peptidoform_index: Some(peptidoform_index),
                            neutral_loss: loss.cloned().unwrap_or_default().into(),
                            deviation: None,
                            confidence: None,
                            auxiliary: false,
                        });
                    }
                }
            }
        }
        result
    }

    /// Generate a list of possible fragments from the list of possible preceding termini and
    /// neutral losses. Ignores any neutral loss that would result in a negative number of any
    /// element.
    /// # Panics
    /// When the charge range results in a negative charge
    #[must_use]
    #[expect(clippy::too_many_arguments)] // Needs many different pieces of information
    pub fn generate_series(
        theoretical_mass: &Multi<Mode::Output>,
        peptidoform_ion_index: usize,
        peptidoform_index: usize,
        annotation: &FragmentType,
        termini: &Multi<Mode::Output>,
        neutral_losses: &[Vec<NeutralLoss>],
        charge_carriers: &mut CachedCharge,
        settings: &PossiblePrimaryIons,
    ) -> Vec<Self> {
        let charges = charge_carriers.range(settings.1);
        let losses = std::iter::once(None)
            .chain(settings.0.iter().map(Some))
            .chain(neutral_losses.iter().map(Some))
            .collect::<Vec<_>>();
        let mut result = Vec::with_capacity(
            termini.len()
                * theoretical_mass.len()
                * charges.len()
                * losses.len()
                * settings.2.len(),
        );

        for term in termini.iter() {
            for mass in theoretical_mass.iter() {
                let f = term.clone() + mass.clone();
                for charge in &charges {
                    let f = f.clone()
                        + charge.calculate_mass_inner::<Mode>(
                            SequencePosition::default(),
                            peptidoform_index,
                        );
                    if f.contains_negative_amount() {
                        continue;
                    }
                    let z = Charge::new::<system::e>(charge.charge().value);
                    for loss in &losses {
                        let f = f.clone()
                            + Mode::from_formula(
                                loss.iter().flat_map(|l| l.iter()).sum::<MolecularFormula>(),
                            );
                        for variant in settings.2 {
                            let f =
                                f.clone() + Mode::from_formula(molecular_formula!(H 1) * variant);
                            if f.contains_negative_amount() {
                                continue;
                            }
                            result.push(Self {
                                formula: Some(f),
                                charge: z,
                                ion: annotation.with_variant(*variant),
                                isotope: ThinVec::new(),
                                peptidoform_ion_index: Some(peptidoform_ion_index),
                                peptidoform_index: Some(peptidoform_index),
                                neutral_loss: loss.cloned().unwrap_or_default().into(),
                                deviation: None,
                                confidence: None,
                                auxiliary: false,
                            });
                        }
                    }
                }
            }
        }
        result
    }

    /// Create a copy of this fragment with the given charge
    #[must_use]
    fn with_charge(&self, charge: &MolecularCharge) -> Self {
        let formula =
            charge
                .calculate_mass::<Mode>()
                .with_labels(&[AmbiguousLabel::ChargeCarrier(
                    charge.calculate_mass::<OutputMolecularFormula>(),
                )]);
        let c = Charge::new::<system::charge::e>(formula.charge().value);
        Self {
            formula: Some(self.formula.clone().unwrap_or_default() + formula),
            charge: c,
            ..self.clone()
        }
    }

    /// Create a copy of this fragment with the given charges
    pub fn with_charge_range(
        self,
        charge_carriers: &mut CachedCharge,
        charge_range: ChargeRange,
    ) -> impl Iterator<Item = Self> {
        charge_carriers
            .range(charge_range)
            .into_iter()
            .map(move |c| self.with_charge(&c))
    }

    /// Create a copy of this fragment with the given charges
    pub fn with_charge_range_slice(
        self,
        charges: &[MolecularCharge],
    ) -> impl Iterator<Item = Self> {
        charges.iter().map(move |c| self.with_charge(c))
    }

    /// Create a copy of this fragment with the given neutral loss. This could result in a molecular
    /// formula with an element with a negative amount, use
    /// [`MolecularFormula::contains_negative_amount`] to check for this.
    #[must_use]
    pub fn with_neutral_loss(&self, neutral_loss: &NeutralLoss) -> Self {
        let mut new_neutral_loss = self.neutral_loss.clone();
        new_neutral_loss.push(neutral_loss.clone());
        Self {
            formula: Some(
                self.formula.clone().unwrap_or_default() + neutral_loss.calculate_mass::<Mode>(),
            ),
            neutral_loss: new_neutral_loss,
            ..self.clone()
        }
    }

    /// Create copies of this fragment with the given neutral losses (and a copy of this fragment
    /// itself) ignores any neutral loss that would result in a negative element number.
    #[must_use]
    pub fn with_neutral_losses(&self, neutral_losses: &[NeutralLoss]) -> Vec<Self> {
        let mut output = Vec::with_capacity(neutral_losses.len() + 1);
        output.push(self.clone());
        output.extend(
            neutral_losses
                .iter()
                .map(|loss| self.with_neutral_loss(loss))
                .filter(|f| f.formula.as_ref().is_some_and(|f| !f.contains_negative_amount())),
        );
        output
    }

    /// Update this fragment with this isotope. This also updates the formula. Note that this can
    /// leave the formula with negative amounts of certain elements. If `C2H5O1` gets `+i15N` the
    /// result is `C2H5O1N-1[15N1]`. If this is not acceptable this will have to be checked after
    /// the fact.
    ///
    /// # Panics
    /// If any of the given isotopes are not valid, see [`MolecularFormula::new`].
    #[must_use]
    pub fn with_isotope(mut self, isotopes: &[(i32, Isotope)]) -> Self {
        self.isotope = isotopes.iter().copied().filter(|(a, _)| *a != 0).collect();
        if let Some(formula) = &mut self.formula {
            for (amount, isotope) in &self.isotope {
                isotope.add_to_formula::<Mode>(*amount, formula).unwrap();
            }
        }
        self
    }

    /// Get the base formula with all isotopes removed. It returns None if no formula is available
    /// or if the calculations overflow when doing the subtraction of the isotopes.
    pub fn base_formula(&self) -> Option<Mode::Output> {
        self.formula.clone().and_then(|mut formula| {
            for (amount, isotope) in &self.isotope {
                isotope.sub_from_formula::<Mode>(*amount, &mut formula)?;
            }
            Some(formula)
        })
    }

    /// Convert a formula fragment into this mode (cannot be From because of specialisation)
    pub fn from(value: Fragment<OutputMolecularFormula>) -> Self {
        Self {
            formula: value.formula.map(|f| Mode::from_formula(f)),
            charge: value.charge,
            ion: value.ion,
            isotope: value.isotope,
            peptidoform_ion_index: value.peptidoform_ion_index,
            peptidoform_index: value.peptidoform_index,
            neutral_loss: value.neutral_loss,
            deviation: value.deviation,
            confidence: value.confidence,
            auxiliary: value.auxiliary,
        }
    }
}

impl<Mode: MassOutputMode> Display for Fragment<Mode> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}@{}{:+}{}",
            self.ion,
            self.mz(MassMode::Monoisotopic)
                .map_or(String::new(), |mz| mz.value.to_string()),
            self.charge.value,
            self.neutral_loss.iter().map(ToString::to_string).join("")
        )
    }
}

impl<Mode: MassOutputMode> mzcore::space::Space for Fragment<Mode> {
    fn space(&self) -> mzcore::space::UsedSpace {
        (self.formula.space()
            + self.charge.space()
            + self.ion.space()
            + self.peptidoform_ion_index.space()
            + self.peptidoform_index.space()
            + self.neutral_loss.space()
            + self.deviation.space()
            + self.confidence.space()
            + self.auxiliary.space())
        .set_total::<Self>()
    }
}

/// An isotope type
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
pub enum Isotope {
    /// A general isotope, assumed to be the mass difference between C13 and C12
    General,
    /// An average isotope, the average mass of all isotopologues that make up this isotope
    Average,
    /// A specific isotope
    Specific(Element, NonZeroU16),
}

/// Remember the C13 offset
static ISOTOPE_OFFSET: LazyLock<f64> =
    LazyLock::new(|| molecular_formula!([13 C 1] [12 C -1]).monoisotopic_mass().value);

impl Isotope {
    /// Add a certain amount of this isotope to a formula. It returns None if this isotope is
    /// invalid or if adding this isotope overflows the molecular formula.
    fn add_to_formula<Mode: MassOutputMode>(
        self,
        amount: i32,
        formula: &mut Mode::Output,
    ) -> Option<()> {
        match self {
            Self::Average | Self::General => {
                *formula = formula.clone()
                    + Mode::from_mass(system::Mass::new::<system::dalton>(
                        *ISOTOPE_OFFSET * f64::from(amount),
                    ));
                Some(())
            }
            Self::Specific(el, i) => {
                *formula = formula.clone()
                    + Mode::from_formula(MolecularFormula::new(
                        &[(el, Some(i), amount), (el, None, -amount)],
                        &[],
                    )?);
                Some(())
            }
        }
    }

    /// Subtract a certain amount of this isotope to a formula. It returns None if this isotope is
    /// invalid or if subtracting this isotope overflows the molecular formula.
    fn sub_from_formula<Mode: MassOutputMode>(
        self,
        amount: i32,
        formula: &mut Mode::Output,
    ) -> Option<()> {
        match self {
            Self::Average | Self::General => {
                *formula = formula.clone()
                    + Mode::from_mass(system::Mass::new::<system::dalton>(
                        *ISOTOPE_OFFSET * -1.0 * f64::from(amount),
                    ));
                Some(())
            }
            Self::Specific(el, i) => {
                *formula = formula.clone()
                    - Mode::from_formula(MolecularFormula::new(
                        &[(el, Some(i), amount), (el, None, -amount)],
                        &[],
                    )?);
                Some(())
            }
        }
    }
}

impl Display for Isotope {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::General => write!(f, "i"),
            Self::Average => write!(f, "iA"),
            Self::Specific(el, i) => write!(f, "i{i}{el}"),
        }
    }
}

#[cfg(test)]
#[expect(clippy::missing_panics_doc)]
mod tests {
    use mzcore::sequence::PeptidePosition;

    use super::*;

    #[test]
    fn neutral_loss() {
        let a = Fragment::<OutputMolecularFormula>::new(
            AminoAcid::AsparticAcid.calculate_masses::<OutputMolecularFormula>()[0].clone(),
            Charge::new::<system::charge::e>(1),
            0,
            0,
            FragmentType::Precursor,
        );
        let loss = a.with_neutral_losses(&[NeutralLoss::Loss(1, molecular_formula!(H 2 O 1))]);
        //dbg!(&a, &loss);
        assert_eq!(a.formula, loss[0].formula);
        assert_eq!(
            a.formula.unwrap(),
            &loss[1].formula.clone().unwrap() + &molecular_formula!(H 2 O 1)
        );
    }

    #[test]
    fn flip_terminal() {
        let n0 = PeptidePosition::n(SequencePosition::Index(0, 2), 2);
        let n1 = PeptidePosition::n(SequencePosition::Index(1, 2), 2);
        let n2 = PeptidePosition::n(SequencePosition::Index(2, 2), 2);
        let c0 = PeptidePosition::c(SequencePosition::Index(0, 2), 2);
        let c1 = PeptidePosition::c(SequencePosition::Index(1, 2), 2);
        let c2 = PeptidePosition::c(SequencePosition::Index(2, 2), 2);
        assert_eq!(n0.flip_terminal(), c0);
        assert_eq!(n1.flip_terminal(), c1);
        assert_eq!(n2.flip_terminal(), c2);
    }
}