mzannotate 0.2.0

use std::collections::{HashMap, HashSet};

use itertools::Itertools;
use mzcore::{
    chemistry::{CachedCharge, DiagnosticIon, MassOutputMode, MassOutputType},
    molecular_formula,
    prelude::{MolecularCharge, Peptidoform, PeptidoformIon, PeptidoformIonSet, SequencePosition},
    quantities::Multi,
    sequence::{
        AtMax, GnoComposition, HiddenInternalMethods, Linear, Linked, LinkerSpecificity,
        PeptidePosition, SimpleModificationInner,
    },
    system::isize::Charge,
};
use thin_vec::ThinVec;

use crate::{
    annotation::model::get_all_sidechain_losses,
    fragment::{BackboneCFragment, BackboneNFragment, DiagnosticPosition, FragmentType},
    glycan::GlycanFragmention,
    helper_functions::merge_hashmap,
    modification,
    prelude::{Fragment, FragmentationModel},
};

/// Helper trait to be able to define fragmentation methods on peptidoforms.
pub trait PeptidoformFragmentation {
    /// Generate theoretical fragments with the given maximal charge (ignored if the peptidoform
    /// contains charge carriers) and the given model.
    fn generate_theoretical_fragments<Mode: MassOutputMode>(
        &self,
        max_charge: Charge,
        model: &FragmentationModel,
    ) -> Vec<Fragment<Mode>>;
}

impl PeptidoformFragmentation for PeptidoformIonSet {
    /// Generate the theoretical fragments for this peptidoform ion set.
    fn generate_theoretical_fragments<Mode: MassOutputMode>(
        &self,
        max_charge: Charge,
        model: &FragmentationModel,
    ) -> Vec<Fragment<Mode>> {
        let mut base = Vec::new();
        for (index, peptidoform) in self.peptidoform_ions().iter().enumerate() {
            base.extend(peptidoform_ion_inner(peptidoform, max_charge, model, index));
        }
        base
    }
}

impl PeptidoformFragmentation for PeptidoformIon {
    /// Generate the theoretical fragments for this peptidoform.
    fn generate_theoretical_fragments<Mode: MassOutputMode>(
        &self,
        max_charge: Charge,
        model: &FragmentationModel,
    ) -> Vec<Fragment<Mode>> {
        peptidoform_ion_inner(self, max_charge, model, 0)
    }
}

/// Generate the theoretical fragments for this peptidoform.
fn peptidoform_ion_inner<Mode: MassOutputMode>(
    peptidoform_ion: &PeptidoformIon,
    max_charge: Charge,
    model: &FragmentationModel,
    peptidoform_ion_index: usize,
) -> Vec<Fragment<Mode>> {
    let mut base = Vec::new();
    for (index, peptide) in peptidoform_ion.peptidoforms().iter().enumerate() {
        base.extend(generate_theoretical_fragments_inner(
            peptide,
            max_charge,
            model,
            peptidoform_ion_index,
            index,
            peptidoform_ion.peptidoforms(),
        ));
    }
    base
}

/// Generate the theoretical fragments for this peptide, with the given maximal charge of the
/// fragments, and the given model. With the global isotope modifications applied.
/// # Panics
/// If the global isotope replacement is invalid.
pub(crate) fn generate_theoretical_fragments_inner<Complexity, Mode: MassOutputMode>(
    peptidoform: &Peptidoform<Complexity>,
    max_charge: Charge,
    model: &FragmentationModel,
    peptidoform_ion_index: usize,
    peptidoform_index: usize,
    all_peptides: &[Peptidoform<Linked>],
) -> Vec<Fragment<Mode>> {
    let default_charge = MolecularCharge::proton(max_charge);
    let mut charge_carriers: CachedCharge =
        peptidoform.get_charge_carriers().unwrap_or(&default_charge).into();

    let mut output: Vec<Fragment<Mode>> =
        Vec::with_capacity(20 * peptidoform.sequence().len() + 75); // Empirically derived required size of the buffer (Derived from Hecklib)
    for sequence_index in 0..peptidoform.sequence().len() {
        let position = PeptidePosition::n(
            SequencePosition::Index(sequence_index, peptidoform.len()),
            peptidoform.len(),
        );
        let mut cross_links = Vec::new();
        let visited_peptides = vec![peptidoform_index];
        let (n_term, n_term_specific, n_term_seen, n_term_losses) = peptidoform.all_masses::<Mode>(
            ..=sequence_index,
            ..sequence_index,
            &peptidoform.get_n_term_mass::<Mode>(
                all_peptides,
                &visited_peptides,
                &mut cross_links,
                model.allow_cross_link_cleavage,
                peptidoform_index,
                peptidoform_ion_index,
                &model.glycan,
            ),
            all_peptides,
            &visited_peptides,
            &mut cross_links,
            model.allow_cross_link_cleavage,
            peptidoform_index,
            peptidoform_ion_index,
            &model.glycan,
        );
        let (c_term, c_term_specific, c_term_seen, c_term_losses) = peptidoform.all_masses::<Mode>(
            sequence_index..,
            sequence_index + 1..,
            &peptidoform.get_c_term_mass::<Mode>(
                all_peptides,
                &visited_peptides,
                &mut cross_links,
                model.allow_cross_link_cleavage,
                peptidoform_index,
                peptidoform_ion_index,
                &model.glycan,
            ),
            all_peptides,
            &visited_peptides,
            &mut cross_links,
            model.allow_cross_link_cleavage,
            peptidoform_index,
            peptidoform_ion_index,
            &model.glycan,
        );
        if !n_term_seen.is_disjoint(&c_term_seen) {
            continue; // There is a link reachable from both sides so there is a loop
        }
        let (modifications_total, modifications_specific, modifications_cross_links) =
            peptidoform.sequence()[sequence_index].modifications.iter().fold(
                (Multi::default(), HashMap::new(), HashSet::new()),
                |acc, m| {
                    let (f, specific, s) = m.formula_inner::<Mode>(
                        all_peptides,
                        &[peptidoform_index],
                        &mut cross_links,
                        model.allow_cross_link_cleavage,
                        SequencePosition::Index(sequence_index, peptidoform.len()),
                        peptidoform_index,
                        peptidoform_ion_index,
                        &model.glycan,
                        Some(peptidoform.sequence()[sequence_index].aminoacid.aminoacid()),
                    );
                    (
                        acc.0 * f,
                        merge_hashmap(acc.1, specific),
                        acc.2.union(&s).cloned().collect(),
                    )
                },
            );

        output.append(&mut crate::aminoacid::fragments(
            peptidoform.sequence()[sequence_index].aminoacid.aminoacid(),
            &(n_term, n_term_specific, n_term_losses),
            &(c_term, c_term_specific, c_term_losses),
            &(modifications_total, modifications_specific),
            &mut charge_carriers,
            SequencePosition::Index(sequence_index, peptidoform.len()),
            peptidoform.sequence().len(),
            &model.ions(position, peptidoform),
            peptidoform_ion_index,
            peptidoform_index,
            (
                // Allow any N terminal fragment if there is no cross-link to the C terminal side
                c_term_seen.is_disjoint(&modifications_cross_links),
                n_term_seen.is_disjoint(&modifications_cross_links),
            ),
        ));
    }

    // Internal fragments
    if let Some((
        internal_range,
        neutral_losses,
        specific_losses,
        side_chain_losses,
        charge_range,
    )) = &model.internal
    {
        let options = (0..peptidoform.len())
            .map(|i| {
                let o = model.ions(
                    PeptidePosition::n(
                        SequencePosition::Index(i, peptidoform.len()),
                        peptidoform.len(),
                    ),
                    peptidoform,
                );
                (
                    o.a.is_some(),
                    o.b.is_some(),
                    o.c.is_some(),
                    o.x.is_some(),
                    o.y.is_some(),
                    o.z.is_some(),
                )
            })
            .collect::<Vec<_>>();

        for n in 1..peptidoform.len().saturating_sub(*internal_range.start() + 1) {
            for c in n..(peptidoform.len() - 1).min(*internal_range.end() + 1) {
                let o_n = options[n];
                let o_c = options[c];
                if !(o_c.0 || o_c.1 || o_c.2) && !(o_n.3 || o_n.4 || o_n.5) {
                    continue;
                }

                // Add amino acid specific neutral losses
                let mut internal_neutral_losses = specific_losses
                    .iter()
                    .filter_map(|(rule, losses)| {
                        rule.iter()
                            .any(|aa| {
                                peptidoform.sequence()[n..=c]
                                    .iter()
                                    .any(|seq| seq.aminoacid.aminoacid() == *aa)
                            })
                            .then_some(losses)
                    })
                    .flatten()
                    .map(|l| vec![l.clone()])
                    .collect::<Vec<_>>();
                // Add amino acid side chain losses
                internal_neutral_losses.extend(get_all_sidechain_losses(
                    &peptidoform.sequence()[n..=c],
                    side_chain_losses,
                ));
                // Add all normal neutral losses
                internal_neutral_losses.extend(neutral_losses.iter().map(|l| vec![l.clone()]));
                // TODO: this adds the full peptidoform if there is any cycle to the original
                // peptidoform again.
                let (mass, specific, _seen, losses) = peptidoform.all_masses::<Mode>(
                    n..=c,
                    n..=c,
                    &(Multi::default(), HashMap::new()),
                    all_peptides,
                    &[],
                    &mut Vec::new(),
                    model.allow_cross_link_cleavage,
                    peptidoform_index,
                    peptidoform_ion_index,
                    &model.glycan,
                );
                if model.modification_specific_neutral_losses {
                    internal_neutral_losses.extend(losses);
                }
                for (frag_n, n_possible) in [
                    (BackboneCFragment::x, o_n.3),
                    (BackboneCFragment::y, o_n.4),
                    (BackboneCFragment::z, o_n.5),
                ] {
                    if !n_possible {
                        continue;
                    }
                    for (frag_c, c_possible) in [
                        (BackboneNFragment::a, o_c.0),
                        (BackboneNFragment::b, o_c.1),
                        (BackboneNFragment::c, o_c.2),
                    ] {
                        if c_possible {
                            output.extend(Fragment::generate_all(
                                &((specific.get(&frag_n.into()).unwrap_or(&mass).clone()
                                    * specific.get(&frag_c.into()).unwrap_or(&mass))
                                .unique()
                                    + match (frag_n, frag_c) {
                                        (BackboneCFragment::y, BackboneNFragment::a) => {
                                            Mode::from_formula(molecular_formula!(C -1 O -1))
                                        }
                                        (BackboneCFragment::z, BackboneNFragment::a) => {
                                            Mode::from_formula(
                                                molecular_formula!(C -1 O -1 N -1 H - 1),
                                            )
                                        }
                                        (BackboneCFragment::x, BackboneNFragment::b) => {
                                            Mode::from_formula(molecular_formula!(C 1 O 1))
                                        }
                                        (BackboneCFragment::z, BackboneNFragment::b) => {
                                            Mode::from_formula(molecular_formula!(N -1 H -1))
                                        }
                                        (BackboneCFragment::x, BackboneNFragment::c) => {
                                            Mode::from_formula(molecular_formula!(C 1 O 1 N 1 H 1))
                                        }
                                        (BackboneCFragment::y, BackboneNFragment::c) => {
                                            Mode::from_formula(molecular_formula!(N 1 H 1))
                                        }
                                        _ => Mode::Output::default(),
                                    }),
                                peptidoform_ion_index,
                                peptidoform_index,
                                &FragmentType::Internal(
                                    Some((frag_n, frag_c)),
                                    PeptidePosition::n(
                                        SequencePosition::Index(n, peptidoform.len()),
                                        peptidoform.len(),
                                    ),
                                    PeptidePosition::c(
                                        SequencePosition::Index(c, peptidoform.len()),
                                        peptidoform.len(),
                                    ),
                                ),
                                &Multi::default(),
                                &internal_neutral_losses,
                                &mut charge_carriers,
                                *charge_range,
                            ));
                        }
                    }
                }
            }
        }
    }
    for fragment in &mut output {
        fragment.formula = fragment.formula.take().map(|f| {
            f.with_global_isotope_modifications(peptidoform.get_global())
                .expect("Invalid global isotope modification")
        });
    }

    // Generate precursor peak
    let (full_precursor, _precursor_specific, _all_cross_links) = peptidoform
        .formulas_inner::<Mode>(
            peptidoform_index,
            peptidoform_ion_index,
            all_peptides,
            &[],
            &mut Vec::new(),
            model.allow_cross_link_cleavage,
            &model.glycan,
        );
    // Allow neutral losses from modifications for the precursor
    let mut precursor_neutral_losses = if model.modification_specific_neutral_losses {
        peptidoform
            .potential_neutral_losses(.., all_peptides, peptidoform_index, &mut Vec::new())
            .into_iter()
            .map(|(n, ..)| vec![n])
            .collect()
    } else {
        Vec::new()
    };
    // Add amino acid specific neutral losses
    precursor_neutral_losses.extend(
        model
            .precursor
            .1
            .iter()
            .filter_map(|(rule, losses)| {
                rule.iter()
                    .any(|aa| {
                        peptidoform.sequence().iter().any(|seq| seq.aminoacid.aminoacid() == *aa)
                    })
                    .then_some(losses)
            })
            .flatten()
            .map(|l| vec![l.clone()]),
    );
    // Add amino acid side chain losses
    precursor_neutral_losses.extend(get_all_sidechain_losses(
        peptidoform.sequence(),
        &model.precursor.2,
    ));
    // Add all normal neutral losses
    precursor_neutral_losses.extend(model.precursor.0.iter().map(|l| vec![l.clone()]));

    output.extend(Fragment::generate_all(
        &full_precursor,
        peptidoform_ion_index,
        peptidoform_index,
        &FragmentType::Precursor,
        &Multi::default(),
        &precursor_neutral_losses,
        &mut charge_carriers,
        model.precursor.3,
    ));

    // Add glycan fragmentation to all peptide fragments
    // Assuming that only one glycan can ever fragment at the same time.
    let full_formula = peptidoform
        .formulas_inner::<Mode>(
            peptidoform_index,
            peptidoform_ion_index,
            all_peptides,
            &[],
            &mut Vec::new(),
            model.allow_cross_link_cleavage,
            &model.glycan,
        )
        .0;
    for (sequence_index, position) in peptidoform.sequence().iter().enumerate() {
        let attachment = (
            position.aminoacid.aminoacid(),
            SequencePosition::Index(sequence_index, peptidoform.len()),
        );
        for modification in &position.modifications {
            output.extend(modification::generate_theoretical_fragments(
                modification,
                model,
                peptidoform_ion_index,
                peptidoform_index,
                &mut charge_carriers,
                &full_formula,
                Some(attachment),
            ));
        }
    }

    if let Some(charge) = model.modification_specific_diagnostic_ions {
        // Add all modification diagnostic ions
        for (dia, pos) in diagnostic_ions(peptidoform) {
            output.extend(
                Fragment {
                    formula: Some(Mode::from_formula(dia.0)),
                    charge: Charge::default(),
                    ion: FragmentType::Diagnostic(pos),
                    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,
                }
                .with_charge_range(&mut charge_carriers, charge),
            );
        }
    }

    // Add labile glycan fragments
    for (modification, _) in peptidoform.get_labile() {
        match &**modification {
            SimpleModificationInner::Glycan(composition) => {
                output.extend(crate::monosaccharide::theoretical_fragments(
                    composition,
                    model,
                    peptidoform_ion_index,
                    peptidoform_index,
                    &mut charge_carriers,
                    &full_formula,
                    None,
                ));
            }
            SimpleModificationInner::GlycanStructure(structure)
            | SimpleModificationInner::Gno {
                composition: GnoComposition::Topology(structure),
                ..
            } => {
                output.extend(
                    structure.clone().determine_positions().generate_theoretical_fragments(
                        model,
                        peptidoform_ion_index,
                        peptidoform_index,
                        &mut charge_carriers,
                        &full_formula,
                        None,
                    ),
                );
            }
            _ => (),
        }
    }

    output
}

/// Find all diagnostic ions for this full peptide
fn diagnostic_ions<Complexity>(
    peptidoform: &Peptidoform<Complexity>,
) -> Vec<(DiagnosticIon, DiagnosticPosition)> {
    peptidoform
        .iter(..)
        .flat_map(|(pos, aa)| {
            aa.diagnostic_ions(
                pos.sequence_index,
                peptidoform.get_n_term(),
                peptidoform.get_c_term(),
            )
            .into_iter()
            .map(move |diagnostic| {
                (
                    diagnostic,
                    DiagnosticPosition::Peptide(pos, aa.aminoacid.aminoacid()),
                )
            })
        })
        .chain(
            peptidoform.get_labile().iter().flat_map(
                move |(modification, _)| match &**modification {
                    SimpleModificationInner::Database { specificities, .. } => specificities
                        .iter()
                        .flat_map(|(_, _, diagnostic)| diagnostic)
                        .map(|diagnostic| {
                            (
                                diagnostic.clone(),
                                DiagnosticPosition::Labile(modification.clone().into()),
                            )
                        })
                        .collect::<Vec<_>>(),
                    SimpleModificationInner::Linker { specificities, .. } => specificities
                        .iter()
                        .flat_map(|rule| match rule {
                            LinkerSpecificity::Symmetric { diagnostic, .. }
                            | LinkerSpecificity::Asymmetric { diagnostic, .. } => diagnostic,
                        })
                        .map(|diagnostic| {
                            (
                                diagnostic.clone(),
                                DiagnosticPosition::Labile(modification.clone().into()),
                            )
                        })
                        .collect::<Vec<_>>(),
                    _ => Vec::new(),
                },
            ),
        )
        .unique()
        .collect()
}

impl<Complexity: AtMax<Linear>> PeptidoformFragmentation for Peptidoform<Complexity> {
    /// Generate the theoretical fragments for this peptide, with the given maximal charge of the
    /// fragments, and the given model. With the global isotope modifications applied.
    /// # Panics
    /// If the global isotope replacement is invalid.
    fn generate_theoretical_fragments<Mode: MassOutputMode>(
        &self,
        max_charge: Charge,
        model: &FragmentationModel,
    ) -> Vec<Fragment<Mode>> {
        generate_theoretical_fragments_inner(self, max_charge, model, 0, 0, &[])
    }
}