zoe 0.0.30

//! A module providing a k-mer encoder using three bits per base.
//!
//! See [`ThreeBitKmerEncoder`] for more details. This also provides the type
//! aliases [`ThreeBitKmerSet`] and [`ThreeBitKmerCounter`].

use crate::{
    data::mappings::THREE_BIT_MAPPING,
    kmer::{Kmer, KmerEncoder, KmerError, KmerLen, KmerSet, MaxLenToType, SupportedKmerLen},
    math::{AnyInt, Uint},
    prelude::KmerCounter,
};
use std::hash::{Hash, Hasher, RandomState};

/// A type alias for a [`KmerLen`] struct with the [`ThreeBitKmerEncoder`] as
/// its encoder.
pub(crate) type ThreeBitKmerLen<const MAX_LEN: usize> = KmerLen<MAX_LEN, ThreeBitKmerEncoder<MAX_LEN>>;

/// A type alias for [`MaxLenToType`] with the [`ThreeBitKmerEncoder`] as the
/// encoder.
pub(crate) type ThreeBitMaxLenToType<const MAX_LEN: usize> = MaxLenToType<MAX_LEN, ThreeBitKmerEncoder<MAX_LEN>>;

/// A type alias for a [`KmerSet`] with the [`ThreeBitKmerEncoder`] as its
/// encoder.
///
/// <div class="warning tip">
///
/// **Tip**
///
/// For guidance on picking the appropriate `MAX_LEN`, see [`SupportedKmerLen`].
///
/// </div>
pub type ThreeBitKmerSet<const MAX_LEN: usize, S = RandomState> = KmerSet<MAX_LEN, ThreeBitKmerEncoder<MAX_LEN>, S>;

/// A type alias for a [`KmerCounter`] with the [`ThreeBitKmerEncoder`] as its
/// encoder.
///
/// <div class="warning tip">
///
/// **Tip**
///
/// For guidance on picking the appropriate `MAX_LEN`, see [`SupportedKmerLen`].
///
/// </div>
pub type ThreeBitKmerCounter<const MAX_LEN: usize, S = RandomState> = KmerCounter<MAX_LEN, ThreeBitKmerEncoder<MAX_LEN>, S>;

/// An encoded k-mer using the [`ThreeBitKmerEncoder`]. In most use cases,
/// encoded k-mers need not be handled directly; [`ThreeBitKmerSet`] and
/// [`ThreeBitKmerCounter`] provide many methods for accomplishing common tasks.
/// If no suitable methods are present, then handling encoded k-mers directly
/// and later decoding them may be appropriate.
///
/// To decode a [`ThreeBitKmerEncoder`], call the `decode_kmer` method of the
/// encoder. For the potential case where the encoder is not available,
/// `Display` is implemented for [`ThreeBitEncodedKmer`], which will output the
/// decoded k-mer. However, it is more performant to decode first.
///
/// Note that the ordering given by [`Ord`] and [`PartialOrd`] is different
/// between [`ThreeBitEncodedKmer`] and the corresponding decoded [`Kmer`].
///
/// [`ThreeBitKmerSet`]: crate::kmer::encoders::three_bit::ThreeBitKmerSet
/// [`ThreeBitKmerCounter`]:
///     crate::kmer::encoders::three_bit::ThreeBitKmerCounter
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
#[repr(transparent)]
pub struct ThreeBitEncodedKmer<const MAX_LEN: usize>(ThreeBitMaxLenToType<MAX_LEN>)
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen;

impl<const MAX_LEN: usize, T: Uint> From<T> for ThreeBitEncodedKmer<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    #[inline]
    fn from(value: T) -> Self {
        Self(value)
    }
}

impl<const MAX_LEN: usize, T: Uint> std::fmt::Display for ThreeBitEncodedKmer<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut kmer = self.0;
        let mut buffer = [0; MAX_LEN];
        let mut start = MAX_LEN;
        while kmer != T::ZERO && start > 0 {
            start -= 1;
            let encoded_base = kmer & T::from_literal(0b111);
            buffer[start] = ThreeBitKmerEncoder::decode_base(encoded_base);
            kmer >>= 3;
        }
        // Safety: The buffer only contains valid ASCII because it is
        // initialized with 0 and ThreeBitKmerEncoder::decode_base always
        // returns a char in b"000NACGT"
        f.write_str(unsafe { std::str::from_utf8_unchecked(&buffer[start..]) })
    }
}

impl<const MAX_LEN: usize, T: Uint> std::fmt::Binary for ThreeBitEncodedKmer<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    #[inline]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Binary::fmt(&self.0, f)
    }
}

/// A [`KmerEncoder`] using three bits to represent each base.
///
/// This allows for `A`, `C`, `G`, `T`, and `N` to all be represented. This
/// encoder does not preserve case or the distinction between `T` and `U`. `N`
/// is used as a catch-all for bases that are not `ACGTUNacgtun`.
///
/// <div class="warning tip">
///
/// **Tip**
///
/// For guidance on picking the appropriate `MAX_LEN`, see [`SupportedKmerLen`].
///
/// </div>
///
/// ## Acknowledgements
///
/// This k-mer encoding is inspired by the 2-bit encoding found in
/// [BBDuk](https://jgi.doe.gov/data-and-tools/software-tools/bbtools/bb-tools-user-guide/bbduk-guide/)
/// from the [BBMap](https://sourceforge.net/projects/bbmap/) project.
#[derive(Copy, Clone, Debug)]
pub struct ThreeBitKmerEncoder<const MAX_LEN: usize>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen, {
    kmer_length: usize,
    kmer_mask:   ThreeBitMaxLenToType<MAX_LEN>,
}

impl<T: Uint, const MAX_LEN: usize> ThreeBitKmerEncoder<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    /// Encodes a single base.
    ///
    /// Five unique bases are represented by the encoder: `Aa`, `Cc`, `Gg`,
    /// `TUtu`, or `N` as a catch-all for any other input.
    #[inline]
    #[must_use]
    pub fn encode_base(base: u8) -> T {
        T::from(THREE_BIT_MAPPING[base])
    }

    /// Decodes a single base.
    ///
    /// The base must have been generated using [`ThreeBitKmerEncoder`],
    /// otherwise this function may panic or have unexpected behavior.
    #[inline]
    #[must_use]
    pub fn decode_base(encoded_base: T) -> u8 {
        // as_usize is valid since encoded_base will be in `3..=8`
        b"000NACGT"[encoded_base.cast_as::<usize>()]
    }
}

impl<const MAX_LEN: usize> PartialEq for ThreeBitKmerEncoder<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.kmer_length == other.kmer_length
    }
}

impl<const MAX_LEN: usize> Eq for ThreeBitKmerEncoder<MAX_LEN> where ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen {}

impl<const MAX_LEN: usize> Hash for ThreeBitKmerEncoder<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.kmer_length.hash(state);
    }
}

impl<const MAX_LEN: usize, T: Uint> KmerEncoder<MAX_LEN> for ThreeBitKmerEncoder<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    type EncodedKmer = ThreeBitEncodedKmer<MAX_LEN>;
    type SeqIter<'a> = ThreeBitKmerIterator<'a, MAX_LEN>;
    type SeqIntoIter = ThreeBitKmerIntoIterator<MAX_LEN>;
    type SeqIterRev<'a> = ThreeBitKmerIteratorRev<'a, MAX_LEN>;
    type SeqIntoIterRev = ThreeBitKmerIntoIteratorRev<MAX_LEN>;

    /// Creates a new [`ThreeBitKmerEncoder`] with the specified k-mer length.
    ///
    /// ## Errors
    ///
    /// Returns [`KmerError::InvalidLength`] if `kmer_length` is less than 2 or
    /// greater than the specified `MAX_LEN`.
    #[inline]
    fn new(kmer_length: usize) -> Result<Self, KmerError> {
        if kmer_length <= MAX_LEN && kmer_length > 1 {
            Ok(Self {
                kmer_length,
                kmer_mask: (T::ONE << (3 * kmer_length)) - T::ONE,
            })
        } else {
            Err(KmerError::InvalidLength)
        }
    }

    #[inline]
    fn kmer_length(&self) -> usize {
        self.kmer_length
    }

    #[inline]
    fn encode_kmer(&self, kmer: impl AsRef<[u8]>) -> Self::EncodedKmer {
        let mut encoded_kmer = T::ZERO;

        for &base in kmer.as_ref() {
            encoded_kmer = (encoded_kmer << 3) | Self::encode_base(base);
        }

        ThreeBitEncodedKmer(encoded_kmer)
    }

    fn decode_kmer(&self, mut encoded_kmer: Self::EncodedKmer) -> Kmer<MAX_LEN> {
        let mut buffer = [0; MAX_LEN];
        for i in (0..self.kmer_length).rev() {
            let encoded_base = encoded_kmer.0 & T::from_literal(0b111);
            buffer[i] = Self::decode_base(encoded_base);
            encoded_kmer.0 >>= 3;
        }
        // Safety: The buffer only contains valid ASCII because it is
        // initialized with 0 and ThreeBitKmerEncoder::decode_base always
        // returns a char in b"000NACGT"
        unsafe { Kmer::new_unchecked(self.kmer_length, buffer) }
    }

    #[inline]
    fn iter_from_sequence<'a, S: AsRef<[u8]> + ?Sized>(&self, seq: &'a S) -> Self::SeqIter<'a> {
        ThreeBitKmerIterator::new(self, seq.as_ref())
    }

    #[inline]
    fn iter_consuming_seq<S: Into<Vec<u8>>>(&self, seq: S) -> Self::SeqIntoIter {
        ThreeBitKmerIntoIterator::new(self, seq.into())
    }

    #[inline]
    fn iter_from_sequence_rev<'a, S: AsRef<[u8]> + ?Sized>(&self, seq: &'a S) -> Self::SeqIterRev<'a> {
        ThreeBitKmerIteratorRev::new(self, seq.as_ref())
    }

    #[inline]
    fn iter_consuming_seq_rev<'a, S: Into<Vec<u8>>>(&self, seq: S) -> Self::SeqIntoIterRev {
        ThreeBitKmerIntoIteratorRev::new(self, seq.into())
    }
}

/// An iterator over the three-bit encoded overlapping k-mers in a sequence,
/// from left to right.
pub struct ThreeBitKmerIterator<'a, const MAX_LEN: usize>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen, {
    current_kmer: ThreeBitEncodedKmer<MAX_LEN>,
    remaining:    std::slice::Iter<'a, u8>,
    kmer_mask:    ThreeBitMaxLenToType<MAX_LEN>,
}

impl<'a, const MAX_LEN: usize> ThreeBitKmerIterator<'a, MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    /// Creates a new [`ThreeBitKmerIterator`].
    #[inline]
    fn new(encoder: &ThreeBitKmerEncoder<MAX_LEN>, seq: &'a [u8]) -> Self {
        if seq.len() < encoder.kmer_length {
            ThreeBitKmerIterator {
                current_kmer: ThreeBitEncodedKmer(ThreeBitMaxLenToType::<MAX_LEN>::ZERO),
                remaining:    [].iter(),
                kmer_mask:    encoder.kmer_mask,
            }
        } else {
            let index = encoder.kmer_length - 1;
            // Validity: This code does not conform to the assumptions of the kmer
            // API, since it uses `encode_kmer` on a kmer that is one base too
            // short. However, each iteration of the loop adds another base as
            // the first operation, so this is corrected when the iterator
            // starts.
            let pre_first_kmer = encoder.encode_kmer(&seq[..index]);

            ThreeBitKmerIterator {
                current_kmer: pre_first_kmer,
                remaining:    seq[index..].iter(),
                kmer_mask:    encoder.kmer_mask,
            }
        }
    }

    /// Helper function to shift a kmer to the left and add a new base. This
    /// encapsulates shared behavior between [`ThreeBitKmerIterator`] and
    /// [`ThreeBitKmerIntoIterator`].
    #[inline]
    #[must_use]
    pub(crate) fn shift_and_add_base(
        new_base: u8, kmer: ThreeBitEncodedKmer<MAX_LEN>, kmer_mask: ThreeBitMaxLenToType<MAX_LEN>,
    ) -> ThreeBitEncodedKmer<MAX_LEN> {
        let encoded_base = ThreeBitKmerEncoder::encode_base(new_base);
        // Shift the k-mer to the left to make room for the new base
        let shifted_kmer = kmer.0 << 3;
        // Add the new base and remove the leftmost base (to preserve the
        // length of the k-mer)
        ((shifted_kmer | encoded_base) & kmer_mask).into()
    }
}

impl<const MAX_LEN: usize> Iterator for ThreeBitKmerIterator<'_, MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    type Item = ThreeBitEncodedKmer<MAX_LEN>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let base = *self.remaining.next()?;
        self.current_kmer = ThreeBitKmerIterator::shift_and_add_base(base, self.current_kmer, self.kmer_mask);
        Some(self.current_kmer)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.remaining.size_hint()
    }
}

impl<const MAX_LEN: usize> ExactSizeIterator for ThreeBitKmerIterator<'_, MAX_LEN> where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen
{
}

/// An iterator over the three-bit encoded overlapping k-mers in a sequence,
/// from left to right, which consumes/stores the sequence. For a non-consuming
/// version, use [`ThreeBitKmerIterator`].
pub struct ThreeBitKmerIntoIterator<const MAX_LEN: usize>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen, {
    seq:          Vec<u8>,
    index:        usize,
    current_kmer: ThreeBitEncodedKmer<MAX_LEN>,
    kmer_mask:    ThreeBitMaxLenToType<MAX_LEN>,
}

impl<const MAX_LEN: usize> ThreeBitKmerIntoIterator<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    /// Creates a new [`ThreeBitKmerIterator`].
    #[inline]
    fn new(encoder: &ThreeBitKmerEncoder<MAX_LEN>, seq: Vec<u8>) -> Self {
        if seq.len() < encoder.kmer_length {
            ThreeBitKmerIntoIterator {
                index: seq.len(),
                seq,
                current_kmer: ThreeBitEncodedKmer(ThreeBitMaxLenToType::<MAX_LEN>::ZERO),
                kmer_mask: encoder.kmer_mask,
            }
        } else {
            let index = encoder.kmer_length - 1;
            // Validity: This code does not conform to the assumptions of the kmer
            // API, since it uses `encode_kmer` on a kmer that is one base too
            // short. However, each iteration of the loop adds another base as
            // the first operation, so this is corrected when the iterator
            // starts.
            let pre_first_kmer = encoder.encode_kmer(&seq[..index]);

            ThreeBitKmerIntoIterator {
                seq,
                index,
                current_kmer: pre_first_kmer,
                kmer_mask: encoder.kmer_mask,
            }
        }
    }
}

impl<const MAX_LEN: usize> Iterator for ThreeBitKmerIntoIterator<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    type Item = ThreeBitEncodedKmer<MAX_LEN>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        // Extract the next base to add on the right end of the k-mer
        let base = *self.seq.get(self.index)?;
        self.index += 1;
        self.current_kmer = ThreeBitKmerIterator::shift_and_add_base(base, self.current_kmer, self.kmer_mask);
        Some(self.current_kmer)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let size = self.seq.len() - self.index;
        (size, Some(size))
    }
}

impl<const MAX_LEN: usize> ExactSizeIterator for ThreeBitKmerIntoIterator<MAX_LEN> where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen
{
}

/// An iterator over the three-bit encoded overlapping k-mers in a sequence,
/// from right to left.
pub struct ThreeBitKmerIteratorRev<'a, const MAX_LEN: usize>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen, {
    current_kmer: ThreeBitEncodedKmer<MAX_LEN>,
    remaining:    std::slice::Iter<'a, u8>,
    kmer_length:  usize,
}

impl<'a, const MAX_LEN: usize> ThreeBitKmerIteratorRev<'a, MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    fn new(encoder: &ThreeBitKmerEncoder<MAX_LEN>, seq: &'a [u8]) -> Self {
        if seq.len() < encoder.kmer_length {
            ThreeBitKmerIteratorRev {
                current_kmer: ThreeBitEncodedKmer(ThreeBitMaxLenToType::<MAX_LEN>::ZERO),
                remaining:    [].iter(),
                kmer_length:  encoder.kmer_length,
            }
        } else {
            let index = seq.len() + 1 - encoder.kmer_length;
            // Validity: This code does not conform to the assumptions of the kmer
            // API, since it uses `encode_kmer` on a kmer that is one base too
            // short. However, each iteration of the loop adds another base as
            // the first operation, so this is corrected when the iterator
            // starts.
            let pre_last_kmer = encoder.encode_kmer(&seq[index..]).0 << 3;

            ThreeBitKmerIteratorRev {
                current_kmer: pre_last_kmer.into(),
                remaining:    seq[..index].iter(),
                kmer_length:  encoder.kmer_length,
            }
        }
    }

    /// Helper function to shift a kmer to the right and add a new base. This
    /// encapsulates shared behavior between [`ThreeBitKmerIteratorRev`] and
    /// [`ThreeBitKmerIntoIteratorRev`].
    #[inline]
    #[must_use]
    pub(crate) fn shift_and_add_base_rev(
        new_base: u8, kmer: ThreeBitEncodedKmer<MAX_LEN>, kmer_length: usize,
    ) -> ThreeBitEncodedKmer<MAX_LEN> {
        let encoded_base = ThreeBitKmerEncoder::encode_base(new_base);
        // Shift the new base to the proper position
        let shifted_encoded_base = encoded_base << (3 * (kmer_length - 1));
        // Shift the k-mer to the right to make room for the new base and remove
        // the rightmost base, then add the new base
        ((kmer.0 >> 3) | shifted_encoded_base).into()
    }
}

impl<const MAX_LEN: usize> Iterator for ThreeBitKmerIteratorRev<'_, MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    type Item = ThreeBitEncodedKmer<MAX_LEN>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        // Extract the next base to add on the left end of the k-mer
        let base = *self.remaining.next_back()?;
        self.current_kmer = ThreeBitKmerIteratorRev::shift_and_add_base_rev(base, self.current_kmer, self.kmer_length);
        Some(self.current_kmer)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.remaining.size_hint()
    }
}

impl<const MAX_LEN: usize> ExactSizeIterator for ThreeBitKmerIteratorRev<'_, MAX_LEN> where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen
{
}

/// An iterator over the three-bit encoded overlapping k-mers in a sequence,
/// from right to left, which consumes/stores the sequence. For a non-consuming
/// version, use [`ThreeBitKmerIteratorRev`].
pub struct ThreeBitKmerIntoIteratorRev<const MAX_LEN: usize>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen, {
    seq:          Vec<u8>,
    index:        usize,
    current_kmer: ThreeBitEncodedKmer<MAX_LEN>,
    kmer_length:  usize,
}

impl<const MAX_LEN: usize> ThreeBitKmerIntoIteratorRev<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    fn new(encoder: &ThreeBitKmerEncoder<MAX_LEN>, seq: Vec<u8>) -> Self {
        if seq.len() < encoder.kmer_length {
            ThreeBitKmerIntoIteratorRev {
                index: 0,
                seq,
                current_kmer: ThreeBitEncodedKmer(ThreeBitMaxLenToType::<MAX_LEN>::ZERO),
                kmer_length: encoder.kmer_length,
            }
        } else {
            let index = seq.len() + 1 - encoder.kmer_length;
            // Validity: This code does not conform to the assumptions of the kmer
            // API, since it uses `encode_kmer` on a kmer that is one base too
            // short. However, each iteration of the loop adds another base as
            // the first operation, so this is corrected when the iterator
            // starts.
            let pre_last_kmer = encoder.encode_kmer(&seq[index..]).0 << 3;

            ThreeBitKmerIntoIteratorRev {
                index,
                seq,
                current_kmer: pre_last_kmer.into(),
                kmer_length: encoder.kmer_length,
            }
        }
    }
}

impl<const MAX_LEN: usize> Iterator for ThreeBitKmerIntoIteratorRev<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    type Item = ThreeBitEncodedKmer<MAX_LEN>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.index = self.index.checked_sub(1)?;
        // Extract the next base to add on the left end of the k-mer
        let base = self.seq[self.index];
        self.current_kmer = ThreeBitKmerIteratorRev::shift_and_add_base_rev(base, self.current_kmer, self.kmer_length);
        Some(self.current_kmer)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.index, Some(self.index))
    }
}

impl<const MAX_LEN: usize> ExactSizeIterator for ThreeBitKmerIntoIteratorRev<MAX_LEN> where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen
{
}

/// An iterator over all three-bit encoded k-mers that are at most a Hamming
/// distance of 1 away from a provided k-mer. The original k-mer is included in
/// the iterator.
pub struct ThreeBitOneMismatchIter<const MAX_LEN: usize>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen, {
    encoded_kmer:       ThreeBitEncodedKmer<MAX_LEN>,
    kmer_length:        usize,
    current_kmer:       ThreeBitEncodedKmer<MAX_LEN>,
    current_index:      usize,
    current_base_num:   usize,
    set_mask_third_bit: ThreeBitMaxLenToType<MAX_LEN>,
    not_finished:       bool,
}

impl<const MAX_LEN: usize, T: Uint> ThreeBitOneMismatchIter<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    /// Create a new [`ThreeBitOneMismatchIter`] from a provided `encoded_kmer`
    /// and `kmer_length`. For most purposes, use [`KmerEncoder::get_variants`]
    /// to obtain an iterator.
    #[inline]
    #[must_use]
    pub(crate) fn new(encoded_kmer: ThreeBitEncodedKmer<MAX_LEN>, kmer_encoder: &ThreeBitKmerEncoder<MAX_LEN>) -> Self {
        Self {
            encoded_kmer,
            kmer_length: kmer_encoder.kmer_length(),
            current_kmer: encoded_kmer,
            current_index: 0,
            current_base_num: 0,
            set_mask_third_bit: T::from_literal(0b100),
            not_finished: true,
        }
    }
}

impl<const MAX_LEN: usize> Iterator for ThreeBitOneMismatchIter<MAX_LEN>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen,
{
    type Item = ThreeBitEncodedKmer<MAX_LEN>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        // Check for end of iterator
        while self.current_index < self.kmer_length {
            // Mutate 4 times to other bases
            if self.current_base_num < 4 {
                self.current_base_num += 1;
                // When current_kmer = 100, then the second summand will be 0
                // and the new value will be 000 due to the first summand
                // getting masked. When current_kmer = 0**, then the first
                // summand is current_kmer and the second summand is 001.
                self.current_kmer = ((self.current_kmer.0 | self.set_mask_third_bit)
                    - ((self.current_kmer.0 & self.set_mask_third_bit) >> 2))
                    .into();
                return Some(self.current_kmer);
            }

            self.current_kmer = self.encoded_kmer;
            self.current_index += 1;
            self.current_base_num = 0;
            self.set_mask_third_bit <<= 3;
        }

        if self.not_finished {
            self.not_finished = false;
            return Some(self.encoded_kmer);
        }
        None
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let size = usize::from(self.not_finished) + 4 * (self.kmer_length - self.current_index) - self.current_base_num;
        (size, Some(size))
    }
}

impl<const MAX_LEN: usize> ExactSizeIterator for ThreeBitOneMismatchIter<MAX_LEN> where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen
{
}

/// An iterator over all three-bit encoded k-mers that are at most a Hamming
/// distance of `N` away from a provided k-mer, where `N >= 2`. The original
/// k-mer is included in the iterator.
///
/// ## Acknowledgements
///
/// The code for subset iteration was inspired by [this
/// blog](https://fishi.devtail.io/weblog/2015/05/18/common-bitwise-techniques-subset-iterations/)
pub struct ThreeBitMismatchIter<const MAX_LEN: usize, const N: usize>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen, {
    /// The original encoded k-mer
    encoded_kmer:            ThreeBitMaxLenToType<MAX_LEN>,
    /// The current encoded k-mer, which is mutated as the iterator cycles
    /// through variants
    current_kmer:            ThreeBitMaxLenToType<MAX_LEN>,
    /// The current subset of bases being mutated, encoded as a bitmask where
    /// ith bit from the left is 1 precisely when the ith base in the kmer is
    /// being mutated
    current_subset:          ThreeBitMaxLenToType<MAX_LEN>,
    /// The number of bases being mutated in `current_subset`
    subset_size:             usize,
    /// A buffer to hold the `set_mask_third_bit` values for each index in
    /// `current_subset`, as well as the number of times each base in the subset
    /// was mutated. This value will always be at least 1, since every base in
    /// the subset must be mutated
    masks_and_times_mutated: [(ThreeBitMaxLenToType<MAX_LEN>, usize); N],
    /// The upper bound for `current_subset`. We must have `current_subset`
    /// strictly less than `max_subset`
    max_subset:              ThreeBitMaxLenToType<MAX_LEN>,
}

impl<const MAX_LEN: usize, const N: usize, T: Uint> ThreeBitMismatchIter<MAX_LEN, N>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    #[inline]
    #[must_use]
    pub(crate) fn new(encoded_kmer: ThreeBitEncodedKmer<MAX_LEN>, encoder: &ThreeBitKmerEncoder<MAX_LEN>) -> Self {
        // This implementation assumes N >= 2
        const { assert!(N >= 2) }

        // Due to the value in the initialization of times_mutated being
        // usize::MAX, this will trigger the creation of a new subset (encoded
        // subset 1) on the first call to next. We could modify the
        // initialization here to avoid this, but weirdly it results in a
        // worsening of performance.
        Self {
            encoded_kmer:            encoded_kmer.0,
            current_kmer:            encoded_kmer.0,
            current_subset:          T::ZERO,
            subset_size:             0,
            masks_and_times_mutated: [(T::from_literal(0b100), usize::MAX); N],
            max_subset:              T::ONE << encoder.kmer_length(),
        }
    }
}

impl<const MAX_LEN: usize, const N: usize, T: Uint> Iterator for ThreeBitMismatchIter<MAX_LEN, N>
where
    ThreeBitKmerLen<MAX_LEN>: SupportedKmerLen<T = T>,
{
    type Item = ThreeBitEncodedKmer<MAX_LEN>;

    fn next(&mut self) -> Option<Self::Item> {
        // Save the output: this iterator yields then updates
        let out = ThreeBitEncodedKmer(self.current_kmer);

        if self.masks_and_times_mutated[0].1 < 4 {
            // This first branch is redundant with the next, but increases
            // efficiency. Mutate the first base in the subset.

            self.current_kmer = (self.current_kmer | self.masks_and_times_mutated[0].0)
                - ((self.current_kmer & self.masks_and_times_mutated[0].0) >> 2);
            self.masks_and_times_mutated[0].1 += 1;
        } else if let Some(i) = self.masks_and_times_mutated[..self.subset_size]
            .iter()
            .position(|(_, base_num)| *base_num < 4)
        {
            // This second branch means that a base was found which hasn't been
            // modified 4 times yet (namely the leftmost).

            // Reset each base before it
            for (set_mask_third_bit, base_num) in &mut self.masks_and_times_mutated[..i] {
                // Modify twice in order to cycle to original base, then first
                // variant base. This may seem inefficient, but it is cheaper
                // than storing a stack of encoded kmers.
                let set_mask_third_bit = *set_mask_third_bit;
                self.current_kmer =
                    (self.current_kmer | set_mask_third_bit) - ((self.current_kmer & set_mask_third_bit) >> 2);
                self.current_kmer =
                    (self.current_kmer | set_mask_third_bit) - ((self.current_kmer & set_mask_third_bit) >> 2);
                *base_num = 1;
            }

            // Mutate the identified base once
            let set_mask_third_bit = self.masks_and_times_mutated[i].0;
            self.current_kmer = (self.current_kmer | set_mask_third_bit) - ((self.current_kmer & set_mask_third_bit) >> 2);
            self.masks_and_times_mutated[i].1 += 1;
        } else {
            // This final branch means we need a new subset. In either case of
            // the below code, the encoded subset (represented as T) will get
            // larger.

            if (self.current_subset.count_ones() as usize) < N {
                // Our subset has not reached maximum size, so we can add 1
                self.current_subset += T::ONE;
            } else {
                // Replace all trailing 0s with 1s, then add 1. This always
                // decreases the number of 1s
                self.current_subset = (self.current_subset | (self.current_subset - T::ONE)) + T::ONE;
            }

            // Check if we are at the end of the iterator
            if self.current_subset >= self.max_subset {
                // The iterator has just reached its end, but we may still need
                // to return Some(out). The subsets are visited in order, so
                // since N >= 1, we will visit self.max_mask (it has only a
                // single element). Hence, when we first arrive at this point in
                // the code, we will have self.current_subset equal to
                // self.max_mask. If this is the case, the below variable
                // becomes Some(out), otherwise it is None.
                let out = (self.current_subset == self.max_subset).then_some(out);

                // Adjust current_subset to make it larger than max_mask, so
                // that we return None from now on. This will not cause
                // overflow, since we are only utilizing one third of the bits
                // in T.
                self.current_subset = self.max_subset + T::ONE;
                return out;
            }

            // Reset current_kmer
            self.current_kmer = self.encoded_kmer;

            // Reprocess the saved state for this new subset
            self.subset_size = 0;
            let mut processed_bases = 0;
            let mut temp_subset = self.current_subset;
            while temp_subset != T::ZERO {
                let num_zeros = temp_subset.trailing_zeros() as usize;
                temp_subset >>= num_zeros + 1;
                processed_bases += num_zeros;
                self.masks_and_times_mutated[self.subset_size].0 = T::from_literal(0b100) << (processed_bases * 3);
                processed_bases += 1;
                self.subset_size += 1;
            }

            // Precompute first variant: must modify each base in the subset
            for (set_mask_third_bit, base_num) in &mut self.masks_and_times_mutated[..self.subset_size] {
                self.current_kmer =
                    (self.current_kmer | *set_mask_third_bit) - ((self.current_kmer & *set_mask_third_bit) >> 2);
                *base_num = 1;
            }
        }

        Some(out)
    }

    // We could implement size_hint and ExactSizeIterator, but it is expensive
    // to compute the size of the iterator after it has already started
}

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

    #[test]
    fn test_valid_decoding() {
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(0), b'0');
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(1), b'0');
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(2), b'0');
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(3), b'N');
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(4), b'A');
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(5), b'C');
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(6), b'G');
        assert_eq!(ThreeBitKmerEncoder::<21>::decode_base(7), b'T');
    }

    #[test]
    #[should_panic(expected = "index out of bounds: the len is 8 but the index is 8")]
    fn test_invalid_decoding() {
        let _ = ThreeBitKmerEncoder::<21>::decode_base(8);
    }
}