zoe 0.0.31

//! Structs and traits to enable more readable/correct indexing into
//! pHMM-related data structures.

use crate::{
    alignment::phmm::{
        CorePhmm, DomainPhmm, GlobalPhmm, LocalPhmm, SemiLocalPhmm,
        indexing::{GetCore, GetLayer},
        modules::{PrecomputedDomainModule, PrecomputedLocalModule, SemiLocalModule},
    },
    data::views::IndexAdjustable,
};
use std::ops::{Bound, Range, RangeBounds, RangeFrom, RangeInclusive, RangeTo, RangeToInclusive};

/// A trait similar to [`RangeBounds`] but for ranges of [`PhmmIndex`] values.
pub trait PhmmIndexRange: Clone {
    /// The type of the starting index.
    type Start: PhmmIndex;
    /// The type of the ending index.
    type End: PhmmIndex;

    /// Returns the starting index in the range as a [`Bound`].
    fn start_bound(&self) -> Bound<Self::Start>;

    /// Returns the ending index in the range as a [`Bound`].
    fn end_bound(&self) -> Bound<Self::End>;
}

impl<I: PhmmIndex, J: PhmmIndex> PhmmIndexRange for (Bound<I>, Bound<J>) {
    type Start = I;
    type End = J;

    #[inline]
    fn start_bound(&self) -> Bound<Self::Start> {
        self.0
    }

    #[inline]
    fn end_bound(&self) -> Bound<Self::End> {
        self.1
    }
}

impl<I: PhmmIndex> PhmmIndexRange for Range<I> {
    type Start = I;
    type End = I;

    fn start_bound(&self) -> Bound<Self::Start> {
        <Self as RangeBounds<I>>::start_bound(self).map(|x| *x)
    }

    fn end_bound(&self) -> Bound<Self::End> {
        <Self as RangeBounds<I>>::end_bound(self).map(|x| *x)
    }
}

impl<I: PhmmIndex> PhmmIndexRange for RangeInclusive<I> {
    type Start = I;
    type End = I;

    fn start_bound(&self) -> Bound<Self::Start> {
        <Self as RangeBounds<I>>::start_bound(self).map(|x| *x)
    }

    fn end_bound(&self) -> Bound<Self::End> {
        <Self as RangeBounds<I>>::end_bound(self).map(|x| *x)
    }
}

impl<I: PhmmIndex> PhmmIndexRange for RangeFrom<I> {
    type Start = I;
    type End = I;

    fn start_bound(&self) -> Bound<Self::Start> {
        <Self as RangeBounds<I>>::start_bound(self).map(|x| *x)
    }

    fn end_bound(&self) -> Bound<Self::End> {
        <Self as RangeBounds<I>>::end_bound(self).map(|x| *x)
    }
}

impl<I: PhmmIndex> PhmmIndexRange for RangeTo<I> {
    type Start = I;
    type End = I;

    fn start_bound(&self) -> Bound<Self::Start> {
        <Self as RangeBounds<I>>::start_bound(self).map(|x| *x)
    }

    fn end_bound(&self) -> Bound<Self::End> {
        <Self as RangeBounds<I>>::end_bound(self).map(|x| *x)
    }
}

impl<I: PhmmIndex> PhmmIndexRange for RangeToInclusive<I> {
    type Start = I;
    type End = I;

    fn start_bound(&self) -> Bound<Self::Start> {
        <Self as RangeBounds<I>>::start_bound(self).map(|x| *x)
    }

    fn end_bound(&self) -> Bound<Self::End> {
        <Self as RangeBounds<I>>::end_bound(self).map(|x| *x)
    }
}

/// A trait for structures that can be indexed via a [`PhmmIndex`], such as
/// pHMMs and modules.
pub trait PhmmIndexable: Sized {
    /// Returns the number of pseudomatch states in the pHMM-related structure
    /// (either a match state for the reference, or BEGIN or END).
    ///
    /// ## Validity
    ///
    /// This must return at least 2.
    #[must_use]
    fn num_pseudomatch(&self) -> usize;

    /// Returns the length of the reference which the pHMM-related structure
    /// represents.
    ///
    /// This is also the number of match states in the pHMM-related structure,
    /// excluding BEGIN and END.
    #[inline]
    #[must_use]
    fn seq_len(&self) -> usize {
        // Validity: num_pseudomatch returns at least 2
        self.num_pseudomatch() - 2
    }

    /// Returns the [`PhmmIndex`] as a dynamic programming index.
    #[inline]
    #[must_use]
    fn get_dp_index(&self, j: impl PhmmIndex) -> usize {
        j.get_phmm_dp_index(self)
    }

    /// Converts the [`PhmmIndex`] to a dynamic programming index.
    #[inline]
    #[must_use]
    #[allow(dead_code)]
    fn to_dp_index(&self, j: impl PhmmIndex) -> DpIndex {
        DpIndex(self.get_dp_index(j))
    }

    /// Returns the [`PhmmIndex`] as a sequence index (with respect to the
    /// reference represented by the pHMM).
    ///
    /// If the index corresponds to the BEGIN state, then `None` is returned
    /// since this does not correspond to a position in the reference.
    #[inline]
    #[must_use]
    fn get_seq_index(&self, j: impl PhmmIndex) -> Option<usize> {
        self.get_dp_index(j).checked_sub(1)
    }

    /// Converts the [`PhmmIndex`] to a sequence index (with respect to the
    /// reference represented by the pHMM).
    ///
    /// If the index corresponds to the BEGIN state, then `None` is returned
    /// since this does not correspond to a position in the reference.
    #[inline]
    #[must_use]
    fn to_seq_index(&self, j: impl PhmmIndex) -> Option<SeqIndex> {
        self.get_seq_index(j).map(SeqIndex)
    }

    /// Returns a range of dynamic programming indices from a
    /// [`PhmmIndexRange`].
    #[inline]
    #[must_use]
    fn get_dp_range<R: PhmmIndexRange>(&self, range: R) -> Range<usize> {
        let start = match range.start_bound() {
            Bound::Included(start) => self.get_dp_index(start),
            // +1 due to converting Excluded to Included
            Bound::Excluded(start) => self.get_dp_index(start) + 1,
            Bound::Unbounded => 0,
        };
        let end = match range.end_bound() {
            // +1 due to converting Included to Excluded
            Bound::Included(end) => self.get_dp_index(end) + 1,
            Bound::Excluded(end) => self.get_dp_index(end),
            Bound::Unbounded => self.seq_len(),
        };

        start..end
    }

    /// Converts a [`PhmmIndexRange`] to a range of [`DpIndex`].
    #[inline]
    #[must_use]
    #[allow(dead_code)]
    fn to_dp_range<R: PhmmIndexRange>(&self, range: R) -> Range<DpIndex> {
        let Range { start, end } = self.get_dp_range(range);
        DpIndex(start)..DpIndex(end)
    }

    /// Returns a range of sequence indices from [`PhmmIndexRange`].
    ///
    /// If either index corresponds to the BEGIN state, then this will be mapped
    /// to 0 (the same sequence index that [`FirstMatch`] corresponds to).
    #[inline]
    #[must_use]
    fn get_seq_range<R: PhmmIndexRange>(&self, range: R) -> Range<usize> {
        // -1 for converting dynamic programming index to sequence index
        self.get_dp_range(range).saturating_sub(1)
    }

    /// Converts a [`PhmmIndexRange`] to a range of [`SeqIndex`].
    ///
    /// If either index corresponds to the BEGIN state, then this will be mapped
    /// to 0 (the same sequence index that [`FirstMatch`] corresponds to).
    #[inline]
    #[must_use]
    #[allow(dead_code)]
    fn to_seq_range<R: PhmmIndexRange>(&self, range: R) -> Range<SeqIndex> {
        let Range { start, end } = self.get_seq_range(range);
        SeqIndex(start)..SeqIndex(end)
    }
}

/// A trait for structures that can be indexed via a [`QueryIndex`], such as
/// query sequence.
pub trait QueryIndexable: Sized {
    /// Returns the length of the query sequence.
    #[must_use]
    fn seq_len(&self) -> usize;

    /// Returns the number of possible states that could be in when traversing
    /// the query.
    ///
    /// Either none of the bases are consumed, or up to and including
    /// `query_len` bases could be consumed.
    #[inline]
    #[must_use]
    fn dp_len(&self) -> usize {
        self.seq_len() + 1
    }

    /// Returns the [`QueryIndex`] as a dynamic programming index.
    #[inline]
    #[must_use]
    fn get_dp_index(&self, i: impl QueryIndex) -> usize {
        i.get_query_dp_index(self)
    }

    /// Converts the [`QueryIndex`] to a dynamic programming index.
    #[inline]
    #[must_use]
    fn to_dp_index(&self, i: impl QueryIndex) -> DpIndex {
        DpIndex(self.get_dp_index(i))
    }

    /// Returns a range of dynamic programming indices from a start and end
    /// [`QueryIndex`].
    #[inline]
    #[must_use]
    fn get_dp_range(&self, start: Bound<impl QueryIndex>, end: Bound<impl QueryIndex>) -> Range<usize> {
        let start = match start {
            Bound::Included(start) => self.get_dp_index(start),
            Bound::Excluded(start) => self.get_dp_index(start) + 1,
            Bound::Unbounded => 0,
        };
        let end = match end {
            Bound::Included(end) => self.get_dp_index(end) + 1,
            Bound::Excluded(end) => self.get_dp_index(end),
            Bound::Unbounded => self.seq_len(),
        };

        start..end
    }

    /// Returns a range of sequence indices from a start and end [`QueryIndex`].
    ///
    /// If either index corresponds [`NoBases`], then this will be mapped to 0
    /// (the same sequence index that [`FirstBase`] corresponds to).
    #[inline]
    #[must_use]
    fn get_seq_range(&self, start: Bound<impl QueryIndex>, end: Bound<impl QueryIndex>) -> Range<usize> {
        // -1 for converting dynamic programming index to sequence index
        self.get_dp_range(start, end).saturating_sub(1)
    }
}

/// A trait representing different ways to index into a pHMM-related data
/// structure.
///
/// By indexing data structures with a [`PhmmIndex`], it allows for better
/// readability and correctness by requiring specification of the type of index
/// (such as a dynamic programming index with [`DpIndex`] or a sequence index
/// with [`SeqIndex`]). It also allows special elements to be accessed with
/// [`Begin`], [`FirstMatch`], [`LastMatch`], and [`End`].
pub trait PhmmIndex: Copy {
    /// Helper function for [`PhmmIndexable::get_dp_index`], allowing each index
    /// type to control how it gets coverted to a dynamic programming index
    #[must_use]
    fn get_phmm_dp_index(self, v: &impl PhmmIndexable) -> usize;

    /// Gets the index before the current one, as a [`DpIndex`].
    ///
    /// ## Panics
    ///
    /// If the previous index is out of bounds (e.g., for [`Begin`] or
    /// `DpIndex(0)`), then this will panic.
    #[inline]
    #[must_use]
    #[allow(dead_code)]
    fn prev_index(self, phmm: &impl PhmmIndexable) -> DpIndex {
        DpIndex(phmm.get_dp_index(self) - 1)
    }

    /// Gets the index after the current one, as a [`DpIndex`].
    ///
    /// No check is performed for whether this index is in bounds for the given
    /// `phmm`.
    #[inline]
    #[must_use]
    fn next_index(self, phmm: &impl PhmmIndexable) -> DpIndex {
        DpIndex(phmm.get_dp_index(self) + 1)
    }

    /// Gets the minimum of two [`PhmmIndex`] structs as a [`DpIndex`] (the
    /// leftmost in the pHMM).
    ///
    /// No bounds checking is performed.
    #[inline]
    #[must_use]
    #[allow(dead_code)]
    fn min_index(self, other: impl PhmmIndex, phmm: &impl PhmmIndexable) -> DpIndex {
        phmm.to_dp_index(self).min(phmm.to_dp_index(other))
    }

    /// Gets the maximum of two [`PhmmIndex`] structs as a [`DpIndex`] (the
    /// rightmost in the pHMM).
    ///
    /// No bounds checking is performed.
    #[inline]
    #[must_use]
    #[allow(dead_code)]
    fn max_index(self, other: impl PhmmIndex, phmm: &impl PhmmIndexable) -> DpIndex {
        phmm.to_dp_index(self).max(phmm.to_dp_index(other))
    }

    /// Tests two indices for equality by converting them both to [`DpIndex`].
    ///
    /// No bounds checking is performed.
    #[inline]
    #[must_use]
    #[allow(dead_code)]
    fn eq_index(self, other: impl PhmmIndex, phmm: &impl PhmmIndexable) -> bool {
        phmm.to_dp_index(self) == phmm.to_dp_index(other)
    }
}

/// A trait representing different ways to index into a query-related data
/// structure.
///
/// By indexing data structures with a [`QueryIndex`], it allows for better
/// readability and correctness by requiring specification of the type of index
/// (such as a dynamic programming index with [`DpIndex`] or a sequence index
/// with [`SeqIndex`]). It also allows special elements to be accessed with
/// [`NoBases`], [`FirstBase`], and [`LastBase`].
pub trait QueryIndex: Copy {
    /// Helper function for [`QueryIndexable::get_dp_index`], allowing each
    /// index type to control how it gets converted to a dynamic programming
    /// index.
    #[must_use]
    fn get_query_dp_index(self, v: &impl QueryIndexable) -> usize;
}

/// A [`PhmmIndex`] or [`QueryIndex`] representing an index with respect to the
/// sequence, rather than with respect to the dynamic programming tables.
///
/// 0 represents the first position in the sequence (reference or query).
#[repr(transparent)]
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct SeqIndex(pub usize);

/// A [`PhmmIndex`] or [`QueryIndex`] representing an index with respect to the
/// dynamic programming tables, rather than with respect to the sequence itself.
///
/// 1 represents the first position in the sequence (reference or query), while
/// 0 represents matching no bases or the BEGIN state of the pHMM.
#[repr(transparent)]
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct DpIndex(pub usize);

/// A [`PhmmIndex`] representing the BEGIN state of the pHMM, as well as the
/// first layer (which holds the transition probabilities out of the BEGIN
/// state).
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Begin;

/// A [`PhmmIndex`] representing the first match state of the pHMM after BEGIN.
///
/// This corresponds to the first residue in the reference sequences.
#[allow(dead_code)]
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct FirstMatch;

/// A [`PhmmIndex`] representing the last match state of the pHMM before END.
///
/// This corresponds to the last residue in the reference sequences.
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct LastMatch;

// TODO: Doc link
/// A [`PhmmIndex`] representing the END state of the pHMM.
///
/// When used to get a layer from a pHMM with `get_layer`, this is treated as
/// the same thing as [`LastMatch`] since the END state does not have its own
/// layer.
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct End;

/// A [`QueryIndex`] representing not matching any bases from the sequence
/// (dynamic programming index 0).
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
#[allow(dead_code)]
pub struct NoBases;

/// A [`QueryIndex`] representing the first base in the sequence.
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
#[allow(dead_code)]
pub struct FirstBase;

/// A [`QueryIndex`] representing the last base in the sequence.
#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct LastBase;

impl<T, const S: usize> PhmmIndexable for CorePhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        // The END state does not have an index in the CorePhmm, so we add 1
        self.layers().len() + 1
    }
}

impl<T, const S: usize> PhmmIndexable for &CorePhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        // The END state does not have an index in the CorePhmm, so we add 1
        self.layers().len() + 1
    }
}

impl<T, const S: usize> PhmmIndexable for &mut CorePhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        // The END state does not have an index in the CorePhmm, so we add 1
        self.layers().len() + 1
    }
}

impl<T, const S: usize> PhmmIndexable for GlobalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &GlobalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &mut GlobalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for LocalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &LocalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &mut LocalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for SemiLocalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &SemiLocalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &mut SemiLocalPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for DomainPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &DomainPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &mut DomainPhmm<T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.core().num_pseudomatch()
    }
}

impl<T> PhmmIndexable for SemiLocalModule<T> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.0.len()
    }
}

impl<T> PhmmIndexable for &SemiLocalModule<T> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.0.len()
    }
}

impl<T> PhmmIndexable for &mut SemiLocalModule<T> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.0.len()
    }
}

impl<T, const S: usize> PhmmIndexable for PrecomputedLocalModule<'_, T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.semilocal_params.num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &PrecomputedLocalModule<'_, T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.semilocal_params.num_pseudomatch()
    }
}

impl<T, const S: usize> PhmmIndexable for &mut PrecomputedLocalModule<'_, T, S> {
    #[inline]
    fn num_pseudomatch(&self) -> usize {
        self.semilocal_params.num_pseudomatch()
    }
}

impl QueryIndexable for &[u8] {
    fn seq_len(&self) -> usize {
        self.len()
    }
}

impl<T, const S: usize> QueryIndexable for PrecomputedDomainModule<T, S> {
    #[inline]
    fn seq_len(&self) -> usize {
        self.0.len() - 1
    }
}

impl PhmmIndex for DpIndex {
    #[inline]
    fn get_phmm_dp_index(self, _v: &impl PhmmIndexable) -> usize {
        self.0
    }
}

impl PhmmIndex for SeqIndex {
    #[inline]
    fn get_phmm_dp_index(self, _v: &impl PhmmIndexable) -> usize {
        self.0 + 1
    }
}

impl PhmmIndex for Begin {
    #[inline]
    fn get_phmm_dp_index(self, _v: &impl PhmmIndexable) -> usize {
        0
    }
}

impl PhmmIndex for FirstMatch {
    #[inline]
    fn get_phmm_dp_index(self, _v: &impl PhmmIndexable) -> usize {
        1
    }
}

impl PhmmIndex for LastMatch {
    #[inline]
    fn get_phmm_dp_index(self, v: &impl PhmmIndexable) -> usize {
        v.num_pseudomatch() - 2
    }
}

impl PhmmIndex for End {
    #[inline]
    fn get_phmm_dp_index(self, v: &impl PhmmIndexable) -> usize {
        v.num_pseudomatch() - 1
    }
}

impl QueryIndex for NoBases {
    #[inline]
    fn get_query_dp_index(self, _v: &impl QueryIndexable) -> usize {
        0
    }
}

impl QueryIndex for FirstBase {
    #[inline]
    fn get_query_dp_index(self, _v: &impl QueryIndexable) -> usize {
        1
    }
}

impl QueryIndex for SeqIndex {
    #[inline]
    fn get_query_dp_index(self, _v: &impl QueryIndexable) -> usize {
        self.0 + 1
    }
}

impl QueryIndex for DpIndex {
    #[inline]
    fn get_query_dp_index(self, _v: &impl QueryIndexable) -> usize {
        self.0
    }
}

impl QueryIndex for LastBase {
    #[inline]
    fn get_query_dp_index(self, v: &impl QueryIndexable) -> usize {
        // The last value in a slice of length dp_len
        v.dp_len() - 1
    }
}