mzalign 0.2.0

use mzcore::{
    chemistry::MassMode,
    quantities::Multi,
    sequence::{HasPeptidoform, Linear},
    system::Mass,
};
#[cfg(feature = "rayon")]
use rayon::prelude::*;

use crate::{
    AlignScoring, AlignType, Alignment,
    diagonal_array::DiagonalArray,
    mass_alignment::{align_cached, calculate_masses},
};

/// An alignment index to store a collection of sequences to do alignment against with one or more
/// other sequences. This improves performance because it can cache some intermediate results.
#[derive(Debug)]
pub struct AlignIndex<const STEPS: u16, Source> {
    pub(super) sequences: Vec<(Source, DiagonalArray<Multi<Mass>, STEPS>)>,
    pub(super) mode: MassMode,
}

impl<const S: u16, P> Default for AlignIndex<S, P> {
    fn default() -> Self {
        Self {
            sequences: Vec::new(),
            mode: MassMode::Monoisotopic,
        }
    }
}

impl<const STEPS: u16, Source: HasPeptidoform<Linear> + Clone> AlignIndex<STEPS, Source> {
    /// Check if this index is empty (contains no sequences)
    pub fn is_empty(&self) -> bool {
        self.sequences.is_empty()
    }

    /// Get the number of sequences in the index
    pub fn len(&self) -> usize {
        self.sequences.len()
    }

    /// Create a new index for alignments
    pub fn new(sequences: impl IntoIterator<Item = Source>, mode: MassMode) -> Self {
        let sequences = sequences
            .into_iter()
            .map(|p| {
                let masses = calculate_masses::<STEPS>(p.cast_peptidoform(), mode);
                (p, masses)
            })
            .collect::<Vec<_>>();

        Self { sequences, mode }
    }

    /// Align a single peptidoform to the index
    pub fn align_one<'a, Other: HasPeptidoform<Linear> + Clone + 'a>(
        &'a self,
        other: Other,
        scoring: AlignScoring<'a>,
        align_type: AlignType,
    ) -> impl ExactSizeIterator<Item = Alignment<Source, Other>> {
        let other_masses = calculate_masses::<STEPS>(other.cast_peptidoform(), self.mode);

        self.sequences.iter().map(move |(template, template_masses)| {
            align_cached::<STEPS, Source, Other>(
                template.clone(),
                template_masses,
                other.clone(),
                &other_masses,
                scoring,
                align_type,
            )
        })
    }

    /// Align a single peptidoform to the index
    pub fn align_one_specific<'a, Other: HasPeptidoform<Linear> + Clone + 'a>(
        &'a self,
        other: Other,
        scoring: AlignScoring<'a>,
        align_type: impl Fn(&Source) -> AlignType,
    ) -> impl ExactSizeIterator<Item = Alignment<Source, Other>> {
        let other_masses = calculate_masses::<STEPS>(other.cast_peptidoform(), self.mode);

        self.sequences.iter().map(move |(template, template_masses)| {
            align_cached::<STEPS, Source, Other>(
                template.clone(),
                template_masses,
                other.clone(),
                &other_masses,
                scoring,
                align_type(template),
            )
        })
    }

    /// Align a single peptidoform to the index but only for the source sequences that match the
    /// given filter.
    pub fn align_one_filtered<'a, Other: HasPeptidoform<Linear> + Clone + 'a>(
        &'a self,
        other: Other,
        filter: impl Fn(&Source) -> bool + 'a,
        scoring: AlignScoring<'a>,
        align_type: AlignType,
    ) -> impl Iterator<Item = Alignment<Source, Other>> {
        let other_masses = calculate_masses::<STEPS>(other.cast_peptidoform(), self.mode);

        self.sequences.iter().filter(move |s| filter(&s.0)).map(
            move |(template, template_masses)| {
                align_cached::<STEPS, Source, Other>(
                    template.clone(),
                    template_masses,
                    other.clone(),
                    &other_masses,
                    scoring,
                    align_type,
                )
            },
        )
    }

    /// Align multiple peptides to this index
    pub fn align<'a, Other: HasPeptidoform<Linear> + Clone + 'a>(
        &'a self,
        others: impl IntoIterator<Item = Other>,
        scoring: AlignScoring<'a>,
        align_type: AlignType,
    ) -> impl Iterator<Item = impl ExactSizeIterator<Item = Alignment<Source, Other>>> {
        others.into_iter().map(move |o| self.align_one(o, scoring, align_type))
    }
}

impl<const STEPS: u16, Source: HasPeptidoform<Linear> + Clone + Sync + Send>
    AlignIndex<STEPS, Source>
{
    #[cfg(feature = "rayon")]
    /// Align a single peptidoform to the index in parallel
    pub fn par_align_one<'a, Other: HasPeptidoform<Linear> + Clone + 'a + Send + Sync>(
        &'a self,
        other: Other,
        scoring: AlignScoring<'a>,
        align_type: AlignType,
    ) -> impl ParallelIterator<Item = Alignment<Source, Other>> {
        let other_masses = calculate_masses::<STEPS>(other.cast_peptidoform(), self.mode);

        self.sequences.par_iter().map(move |(template, template_masses)| {
            align_cached::<STEPS, Source, Other>(
                template.clone(),
                template_masses,
                other.clone(),
                &other_masses,
                scoring,
                align_type,
            )
        })
    }

    #[cfg(feature = "rayon")]
    /// Align a single peptidoform to the index in parallel but only for the source sequences that
    /// match the given filter.
    pub fn par_align_one_filtered<'a, Other: HasPeptidoform<Linear> + Clone + 'a + Send + Sync>(
        &'a self,
        other: Other,
        filter: impl Fn(&Source) -> bool + Sync + Send + 'a,
        scoring: AlignScoring<'a>,
        align_type: AlignType,
    ) -> impl ParallelIterator<Item = Alignment<Source, Other>> {
        let other_masses = calculate_masses::<STEPS>(other.cast_peptidoform(), self.mode);

        self.sequences.par_iter().filter(move |s| filter(&s.0)).map(
            move |(template, template_masses)| {
                align_cached::<STEPS, Source, Other>(
                    template.clone(),
                    template_masses,
                    other.clone(),
                    &other_masses,
                    scoring,
                    align_type,
                )
            },
        )
    }
}

impl<const STEPS: u16, Source: HasPeptidoform<Linear> + Clone + Sync> AlignIndex<STEPS, Source> {
    #[cfg(feature = "rayon")]
    /// Align multiple peptides to this index in parallel
    pub fn par_align<'a, Other: HasPeptidoform<Linear> + Clone + 'a + Sync + Send>(
        &'a self,
        others: impl IntoParallelIterator<Item = Other>,
        scoring: AlignScoring<'a>,
        align_type: AlignType,
    ) -> impl ParallelIterator<Item = impl ExactSizeIterator<Item = Alignment<Source, Other>>> {
        others.into_par_iter().map(move |o| self.align_one(o, scoring, align_type))
    }
}