#![doc = include_str!("../README.md")]

use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::num::TryFromIntError;

use buffertk::Unpackable;
use prototk::FieldNumber;

pub mod binary_search;
pub mod bit_array;
pub mod bit_vector;
pub mod builder;
pub mod encoder;
pub mod isa;
pub mod psi;
pub mod sa;
pub mod sais;
pub mod sampled;
pub mod sigma;
pub mod wavelet_tree;

use crate::bit_vector::BitVector;
use crate::builder::{Builder, Helper};

/////////////////////////////////////////////// Error //////////////////////////////////////////////

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Error {
    IntoUsize,
    Unparseable,
    TextTooLong,
    BadSearch,
    BadTextOffset,
    BadRecordOffset,
    CouldNotConstructBitVector,
    InvalidEncoder,
    InvalidBitVector,
    InvalidWaveletTree,
    InvalidSigma,
    InvalidSuffixArray,
    InvalidInverseSuffixArray,
    InvalidPsi,
    InvalidDocument,
    BadRank(usize),
    BadSelect(usize),
    BadIndex(usize),
    LogicError(&'static str),
}

impl From<TryFromIntError> for Error {
    fn from(_: TryFromIntError) -> Self {
        Self::IntoUsize
    }
}

////////////////////////////////////////////// Offset //////////////////////////////////////////////

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct TextOffset(pub usize);

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct RecordOffset(pub usize);

///////////////////////////////////////////// Document /////////////////////////////////////////////

pub trait Document {
    type Search: Iterator<Item = TextOffset>;

    fn construct<H: Helper>(
        text: Vec<u32>,
        record_boundaries: Vec<usize>,
        builder: &mut Builder<H>,
    ) -> Result<(), Error>;

    fn len(&self) -> usize;

    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn records(&self) -> usize;

    fn search(&self, needle: &[u32]) -> Result<Self::Search, Error>;
    fn count(&self, needle: &[u32]) -> Result<usize, Error>;

    fn lookup(&self, offset: TextOffset) -> Result<RecordOffset, Error>;
    fn retrieve(&self, record: RecordOffset) -> Result<Vec<u32>, Error>;
    fn offset_of(&self, record: RecordOffset) -> Result<TextOffset, Error>;
}

/////////////////////////////////////// ReferenceDocumentStub //////////////////////////////////////

#[derive(Clone, Debug, Default, prototk_derive::Message)]
struct ReferenceDocumentStub {
    #[prototk(1, uint32)]
    text: Vec<u32>,
    #[prototk(2, uint64)]
    record_boundaries: Vec<u64>,
}

impl From<ReferenceDocument> for ReferenceDocumentStub {
    fn from(rd: ReferenceDocument) -> Self {
        let ReferenceDocument {
            text,
            record_boundaries,
        } = rd;
        let record_boundaries = record_boundaries.into_iter().map(|x| x as u64).collect();
        Self {
            text,
            record_boundaries,
        }
    }
}

impl TryFrom<ReferenceDocumentStub> for ReferenceDocument {
    type Error = Error;

    fn try_from(rds: ReferenceDocumentStub) -> Result<Self, Self::Error> {
        let ReferenceDocumentStub {
            text,
            record_boundaries,
        } = rds;
        if record_boundaries.iter().any(|x| *x > usize::MAX as u64) {
            return Err(Error::IntoUsize);
        }
        let record_boundaries = record_boundaries.into_iter().map(|x| x as usize).collect();
        Ok(ReferenceDocument {
            text,
            record_boundaries,
        })
    }
}

///////////////////////////////////////// ReferenceDocument ////////////////////////////////////////

pub struct ReferenceDocument {
    text: Vec<u32>,
    record_boundaries: Vec<usize>,
}

impl Document for ReferenceDocument {
    type Search = std::vec::IntoIter<TextOffset>;

    fn construct<H: Helper>(
        text: Vec<u32>,
        record_boundaries: Vec<usize>,
        builder: &mut Builder<H>,
    ) -> Result<(), Error> {
        check_record_boundaries(&text, &record_boundaries)?;
        builder.append_vec_u32(FieldNumber::must(1), &text);
        builder.append_vec_usize(FieldNumber::must(2), &record_boundaries);
        Ok(())
    }

    fn len(&self) -> usize {
        self.text.len()
    }

    fn records(&self) -> usize {
        self.record_boundaries.len()
    }

    fn search(&self, needle: &[u32]) -> Result<Self::Search, Error> {
        let offsets: Vec<TextOffset> = if needle.is_empty() {
            (0..self.len()).map(TextOffset).collect()
        } else {
            let mut offsets = vec![];
            for (idx, candidate) in self.text.windows(needle.len()).enumerate() {
                if candidate == needle {
                    offsets.push(TextOffset(idx));
                }
            }
            offsets
        };
        Ok(offsets.into_iter())
    }

    fn count(&self, needle: &[u32]) -> Result<usize, Error> {
        Ok(self.search(needle)?.count())
    }

    fn lookup(&self, offset: TextOffset) -> Result<RecordOffset, Error> {
        let partition_point = self.record_boundaries.partition_point(|b| *b < offset.0);
        if partition_point >= self.record_boundaries.len() && !self.record_boundaries.is_empty() {
            Ok(RecordOffset(self.record_boundaries.len() - 1))
        } else if partition_point >= self.record_boundaries.len() {
            Err(Error::BadTextOffset)
        } else if self.record_boundaries[partition_point] <= offset.0 {
            Ok(RecordOffset(partition_point))
        } else {
            Ok(RecordOffset(partition_point - 1))
        }
    }

    fn retrieve(&self, record: RecordOffset) -> Result<Vec<u32>, Error> {
        if record.0 >= self.record_boundaries.len() {
            return Err(Error::BadRecordOffset);
        }
        let record_start = self.record_boundaries[record.0];
        let record_limit = if record.0 + 1 >= self.record_boundaries.len() {
            self.text.len()
        } else {
            self.record_boundaries[record.0 + 1]
        };
        Ok(Vec::from(&self.text[record_start..record_limit]))
    }

    fn offset_of(&self, record: RecordOffset) -> Result<TextOffset, Error> {
        if record.0 >= self.record_boundaries.len() {
            return Err(Error::BadRecordOffset);
        }
        Ok(TextOffset(self.record_boundaries[record.0]))
    }
}

impl<'a> Unpackable<'a> for ReferenceDocument {
    type Error = Error;

    fn unpack<'b: 'a>(buf: &'b [u8]) -> Result<(Self, &'b [u8]), Self::Error> {
        let (rds, buf) = <ReferenceDocumentStub as Unpackable>::unpack(buf)
            .map_err(|_| Error::InvalidDocument)?;
        Ok((rds.try_into()?, buf))
    }
}

////////////////////////////////////////// PsiDocumentStub /////////////////////////////////////////

#[derive(Clone, Debug, Default, prototk_derive::Message)]
struct PsiDocumentStub<'a> {
    #[prototk(1, bytes)]
    record_boundaries: &'a [u8],
    #[prototk(2, bytes)]
    sigma: &'a [u8],
    #[prototk(3, bytes)]
    sa: &'a [u8],
    #[prototk(4, bytes)]
    isa: &'a [u8],
    #[prototk(5, bytes)]
    psi: &'a [u8],
}

//////////////////////////////////////////// PsiDocument ///////////////////////////////////////////

#[derive(Debug)]
pub struct PsiDocument<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    // Record boundaries, where select(i) returns the offset of the i'th record in text.
    record_boundaries: bit_vector::sparse::BitVector<'a>,

    // Translate indices to characters.
    sigma: sigma::Sigma<'a>,

    // Sorted array of all suffixes of the original text.
    //
    // Each entry sa[i] in the suffix array indicates the offset in the original text at which the
    // i'th prefix begins.  Compactly: i<j => text[sa[i]..] < text[sa[j]..].
    sa: SA,

    // Inverse of `sa`.  This maps a location in the text to its location in the suffix array.
    isa: ISA,

    // Psi is the successor function applied to the suffix array.
    // Using psi, you can traverse the suffix array because sa, isa, psi are all related:
    //
    // psi[idx] = isa[sa[idx] + 1]
    //
    // See [DGA] for a more thorough understanding.
    psi: PSI,
}

impl<SA, ISA, PSI> PsiDocument<'_, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    fn backwards_search(&self, needle: &[u32]) -> Result<(usize, usize), Error> {
        let mut needle = needle.iter().rev();
        let mut range = if let Some(t) = needle.next() {
            self.sigma.sa_range_for(*t)?
        } else {
            // If there's no needle, we should return everything except the artificial end marker.
            return Ok((1usize, self.psi.len() - 1));
        };
        // This performs backwards search.
        for t in needle {
            range = self
                .psi
                .constrain(&self.sigma, self.sigma.sa_range_for(*t)?, range)?;
        }
        Ok(range)
    }
}

impl<SA, ISA, PSI> Document for PsiDocument<'_, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    type Search = std::vec::IntoIter<TextOffset>;

    fn construct<H: Helper>(
        text: Vec<u32>,
        record_boundaries: Vec<usize>,
        builder: &mut Builder<H>,
    ) -> Result<(), Error> {
        check_record_boundaries(&text, &record_boundaries)?;
        // compute the record boundaries
        let sparse_boundaries: Vec<usize> =
            record_boundaries[1..].iter().map(|rb| rb - 1).collect();
        bit_vector::sparse::BitVector::from_indices(
            16,
            text.len(),
            &sparse_boundaries,
            &mut builder.sub(FieldNumber::must(1)),
        )
        .ok_or(Error::InvalidBitVector)?;
        // compute sigma
        let mut sigma_builder = builder.sub(FieldNumber::must(2));
        sigma::Sigma::construct(text.iter().copied(), &mut sigma_builder)?;
        let sigma_buf = sigma_builder.relative_bytes(0).to_vec();
        let sigma = sigma::Sigma::unpack(&sigma_buf)?.0;
        drop(sigma_builder);
        macro_rules! construct_u32_offsets {
            ($s:expr, $sais:path) => {{
                let s = $s;
                // compute the suffix array
                let mut sa = vec![0u32; s.len()];
                $sais(&sigma, &s, &mut sa)?;
                drop(s);
                // TODO(rescrv): parameterize
                SA::construct_u32(6, &sa, &mut builder.sub(FieldNumber::must(3)))?;
                // compute the inverse suffix array
                let (isa, psi) = inverse_and_psi_u32(&sa);
                ISA::construct_u32(
                    &isa,
                    &record_boundaries,
                    &mut builder.sub(FieldNumber::must(4)),
                )?;
                drop(sa);
                drop(isa);
                PSI::construct_u32(&sigma, &psi, &mut builder.sub(FieldNumber::must(5)))?;
            }};
        }
        let fits_u32_offsets = text.len() < u32::MAX as usize - 1;
        if fits_u32_offsets && sigma.K() <= u8::MAX as usize + 1 {
            let s = translate_text_u8(&text, &sigma)?;
            drop(text);
            construct_u32_offsets!(s, sais::sais_u8_u32);
        } else if fits_u32_offsets && sigma.K() <= u16::MAX as usize + 1 {
            let s = translate_text_u16(&text, &sigma)?;
            drop(text);
            construct_u32_offsets!(s, sais::sais_u16_u32);
        } else {
            // convert the text into a compacted alphabet
            let s = translate_text::<u32>(&text, &sigma)?;
            drop(text);
            // compute the suffix array
            if s.len() < u32::MAX as usize {
                let mut sa = vec![0u32; s.len()];
                sais::sais_u32(&sigma, &s, &mut sa)?;
                drop(s);
                // TODO(rescrv): parameterize
                SA::construct_u32(6, &sa, &mut builder.sub(FieldNumber::must(3)))?;
                // compute the inverse suffix array
                let (isa, psi) = inverse_and_psi_u32(&sa);
                ISA::construct_u32(
                    &isa,
                    &record_boundaries,
                    &mut builder.sub(FieldNumber::must(4)),
                )?;
                drop(sa);
                drop(isa);
                PSI::construct_u32(&sigma, &psi, &mut builder.sub(FieldNumber::must(5)))?;
            } else {
                let mut sa = vec![0usize; s.len()];
                sais::sais(&sigma, &s, &mut sa)?;
                drop(s);
                // TODO(rescrv): parameterize
                SA::construct(6, &sa, &mut builder.sub(FieldNumber::must(3)))?;
                // compute the inverse suffix array
                let isa = inverse(&sa);
                drop(sa);
                ISA::construct(
                    &isa,
                    &record_boundaries,
                    &mut builder.sub(FieldNumber::must(4)),
                )?;
                // compute the successor array, psi
                let psi = psi::compute(&isa);
                drop(isa);
                PSI::construct(&sigma, &psi, &mut builder.sub(FieldNumber::must(5)))?;
            }
        }
        Ok(())
    }

    fn len(&self) -> usize {
        // psi always includes an implicit end character that we strip
        self.psi.len() - 1
    }

    fn records(&self) -> usize {
        // TODO(rescrv): better error handling; requires modifying Document trait.
        self.record_boundaries
            .rank(self.record_boundaries.len())
            .unwrap_or(0)
            + 1
    }

    fn search(&self, needle: &[u32]) -> Result<Self::Search, Error> {
        let range = self.backwards_search(needle)?;
        if range.0 > range.1 {
            Ok(vec![].into_iter())
        } else {
            let mut result = Vec::with_capacity(range.1 - range.0 + 1);
            for offset in range.0..=range.1 {
                result.push(TextOffset(self.sa.lookup(
                    &self.sigma,
                    &self.psi,
                    offset,
                )?));
            }
            result.sort();
            Ok(result.into_iter())
        }
    }

    fn count(&self, needle: &[u32]) -> Result<usize, Error> {
        let range = self.backwards_search(needle)?;
        if range.0 > range.1 {
            Ok(0)
        } else {
            Ok(range.1 - range.0 + 1)
        }
    }

    fn lookup(&self, offset: TextOffset) -> Result<RecordOffset, Error> {
        Ok(RecordOffset(
            self.record_boundaries
                .rank(offset.0)
                .ok_or(Error::BadRank(offset.0))?,
        ))
    }

    fn offset_of(&self, record: RecordOffset) -> Result<TextOffset, Error> {
        Ok(TextOffset(
            self.record_boundaries
                .select(record.0)
                .ok_or(Error::BadSelect(record.0))?,
        ))
    }

    fn retrieve(&self, record: RecordOffset) -> Result<Vec<u32>, Error> {
        let start: usize = self
            .record_boundaries
            .select(record.0)
            .ok_or(Error::BadSelect(record.0))?;
        // TODO(rescrv):  This treats an error as if it's end.  Need to change the bit_vector API.
        let limit: usize = self
            .record_boundaries
            .select(record.0 + 1)
            .unwrap_or(self.len());
        if start > limit {
            return Err(Error::BadRecordOffset);
        }
        let mut idx = self.isa.lookup(start)?;
        let mut result = Vec::with_capacity(limit - start);
        for _ in start..limit {
            result.push(self.sigma.sa_index_to_t(idx).ok_or(Error::InvalidSigma)?);
            idx = self.psi.lookup(&self.sigma, idx)?;
        }
        Ok(result)
    }
}

impl<'a, SA, ISA, PSI> Unpackable<'a> for PsiDocument<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray + Unpackable<'a>,
    ISA: isa::InverseSuffixArray + Unpackable<'a>,
    PSI: psi::Psi + Unpackable<'a>,
    Error: From<<SA as Unpackable<'a>>::Error>
        + From<<ISA as Unpackable<'a>>::Error>
        + From<<PSI as Unpackable<'a>>::Error>,
{
    type Error = Error;

    fn unpack<'b: 'a>(buf: &'b [u8]) -> Result<(Self, &'b [u8]), Self::Error> {
        let (stub, buf) =
            <PsiDocumentStub as Unpackable>::unpack(buf).map_err(|_| Error::InvalidBitVector)?;
        let record_boundaries = bit_vector::sparse::BitVector::new(stub.record_boundaries)
            .ok_or(Error::InvalidBitVector)?;
        let sigma = sigma::Sigma::unpack(stub.sigma)?.0;
        let sa = SA::unpack(stub.sa)?.0;
        let isa = ISA::unpack(stub.isa)?.0;
        let psi = PSI::unpack(stub.psi)?.0;
        Ok((
            PsiDocument {
                record_boundaries,
                sigma,
                sa,
                isa,
                psi,
            },
            buf,
        ))
    }
}

pub type CompressedDocument<'a> = PsiDocument<
    'a,
    sa::SampledSuffixArray<'a>,
    isa::SampledInverseSuffixArray<'a>,
    psi::wavelet_tree::WaveletTreePsi<
        'a,
        wavelet_tree::prefix::WaveletTree<'a, encoder::HuffmanEncoder>,
    >,
>;

///////////////////////////////////////////// Correlate ////////////////////////////////////////////

#[derive(Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd)]
pub struct Exemplar {
    count: usize,
    needle: Vec<u32>,
}

impl Exemplar {
    pub fn count(&self) -> usize {
        self.count
    }

    pub fn text(&self) -> &[u32] {
        &self.needle
    }
}

impl From<ExemplarState> for Exemplar {
    fn from(mut es: ExemplarState) -> Self {
        es.needle.reverse();
        Self {
            count: es.cardinality,
            needle: es.needle,
        }
    }
}

#[derive(Debug, Default, Eq, PartialEq)]
struct CorrelateDocState {
    numer: usize,
    range: (usize, usize),
}

#[derive(Debug, Default, Eq, PartialEq)]
enum ExemplarDocs {
    #[default]
    Empty,
    One(usize, CorrelateDocState),
    Many(HashMap<usize, CorrelateDocState>),
}

impl ExemplarDocs {
    fn with_capacity(capacity: usize) -> Self {
        if capacity <= 1 {
            Self::Empty
        } else {
            Self::Many(HashMap::with_capacity(capacity))
        }
    }

    fn len(&self) -> usize {
        match self {
            Self::Empty => 0,
            Self::One(_, _) => 1,
            Self::Many(docs) => docs.len(),
        }
    }

    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn one(&self) -> Option<(usize, &CorrelateDocState)> {
        match self {
            Self::Empty => None,
            Self::One(idx, state) => Some((*idx, state)),
            Self::Many(docs) if docs.len() == 1 => {
                docs.iter().next().map(|(idx, state)| (*idx, state))
            }
            Self::Many(_) => None,
        }
    }

    fn insert(&mut self, idx: usize, state: CorrelateDocState) -> Option<CorrelateDocState> {
        match std::mem::take(self) {
            Self::Empty => {
                *self = Self::One(idx, state);
                None
            }
            Self::One(existing_idx, existing_state) => {
                if existing_idx == idx {
                    *self = Self::One(idx, state);
                    Some(existing_state)
                } else {
                    let mut docs = HashMap::with_capacity(2);
                    docs.insert(existing_idx, existing_state);
                    let previous = docs.insert(idx, state);
                    *self = Self::Many(docs);
                    previous
                }
            }
            Self::Many(mut docs) => {
                let previous = docs.insert(idx, state);
                *self = Self::Many(docs);
                previous
            }
        }
    }

    fn iter(&self) -> ExemplarDocsIter<'_> {
        match self {
            Self::Empty => ExemplarDocsIter::Empty,
            Self::One(idx, state) => ExemplarDocsIter::One(Some((*idx, state))),
            Self::Many(docs) => ExemplarDocsIter::Many(docs.iter()),
        }
    }

    fn iter_mut(&mut self) -> ExemplarDocsIterMut<'_> {
        match self {
            Self::Empty => ExemplarDocsIterMut::Empty,
            Self::One(idx, state) => ExemplarDocsIterMut::One(Some((*idx, state))),
            Self::Many(docs) => ExemplarDocsIterMut::Many(docs.iter_mut()),
        }
    }
}

enum ExemplarDocsIter<'a> {
    Empty,
    One(Option<(usize, &'a CorrelateDocState)>),
    Many(std::collections::hash_map::Iter<'a, usize, CorrelateDocState>),
}

impl<'a> Iterator for ExemplarDocsIter<'a> {
    type Item = (usize, &'a CorrelateDocState);

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Empty => None,
            Self::One(item) => item.take(),
            Self::Many(iter) => iter.next().map(|(idx, state)| (*idx, state)),
        }
    }
}

enum ExemplarDocsIterMut<'a> {
    Empty,
    One(Option<(usize, &'a mut CorrelateDocState)>),
    Many(std::collections::hash_map::IterMut<'a, usize, CorrelateDocState>),
}

impl<'a> Iterator for ExemplarDocsIterMut<'a> {
    type Item = (usize, &'a mut CorrelateDocState);

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Empty => None,
            Self::One(item) => item.take(),
            Self::Many(iter) => iter.next().map(|(idx, state)| (*idx, state)),
        }
    }
}

#[derive(Debug, Default, Eq, PartialEq)]
struct ExemplarState {
    stop: Vec<u32>,
    needle: Vec<u32>,
    cardinality: usize,
    docs: ExemplarDocs,
}

impl ExemplarState {
    fn cardinality(&self) -> usize {
        self.cardinality
    }
}

impl Ord for ExemplarState {
    fn cmp(&self, other: &ExemplarState) -> Ordering {
        if self == other {
            Ordering::Equal
        } else {
            self.cardinality.cmp(&other.cardinality)
        }
    }
}

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

pub fn exemplars<'a, SA, ISA, PSI>(
    docs: &'a [&'a PsiDocument<'a, SA, ISA, PSI>],
    boundaries: &[(u32, u32)],
) -> Exemplars<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    exemplars_with_min_length(docs, boundaries, 0)
}

pub fn exemplars_with_min_length<'a, SA, ISA, PSI>(
    docs: &'a [&'a PsiDocument<'a, SA, ISA, PSI>],
    boundaries: &[(u32, u32)],
    min_length: usize,
) -> Exemplars<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    let docs = ExemplarDoc::from_documents(docs);
    let mut heap = BinaryHeap::new();
    for boundary in boundaries.iter() {
        let mut state = ExemplarState {
            // the starting boundary because we do backwards search.
            stop: vec![boundary.0],
            needle: vec![boundary.1],
            cardinality: 0,
            docs: ExemplarDocs::with_capacity(docs.len()),
        };
        for (idx, doc) in docs.iter().enumerate() {
            let Ok(range) = doc.document.sigma.sa_range_for(boundary.1) else {
                continue;
            };
            if range.0 > range.1 {
                continue;
            }
            let numer = range.1 - range.0 + 1;
            state.cardinality = state.cardinality.saturating_add(numer);
            state.docs.insert(idx, CorrelateDocState { numer, range });
        }
        if !state.docs.is_empty() {
            heap.push(state);
        }
    }
    Exemplars {
        docs,
        heap,
        min_length,
        sigma_ranges: Vec::new(),
    }
}

pub fn exemplars_from_needle<'a, SA, ISA, PSI>(
    docs: &'a [&'a PsiDocument<'a, SA, ISA, PSI>],
    stop: &[u32],
    needle: Vec<u32>,
) -> Exemplars<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    let docs = ExemplarDoc::from_documents(docs);
    let mut heap = BinaryHeap::new();
    let mut state = ExemplarState {
        stop: stop.to_vec(),
        needle: needle.clone(),
        cardinality: 0,
        docs: ExemplarDocs::with_capacity(docs.len()),
    };
    for (idx, doc) in docs.iter().enumerate() {
        let Ok(range) = doc.document.backwards_search(&state.needle) else {
            continue;
        };
        if range.0 > range.1 {
            continue;
        }
        let numer = range.1 - range.0 + 1;
        state.cardinality = state.cardinality.saturating_add(numer);
        state.docs.insert(idx, CorrelateDocState { numer, range });
    }
    if !state.docs.is_empty() {
        heap.push(state);
    }
    Exemplars {
        docs,
        heap,
        min_length: 0,
        sigma_ranges: Vec::new(),
    }
}

#[derive(Clone, Copy)]
struct ExemplarSymbol {
    text: u32,
    range: (usize, usize),
}

struct ExemplarDoc<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    document: &'a PsiDocument<'a, SA, ISA, PSI>,
    symbols: Vec<ExemplarSymbol>,
    symbols_by_sigma: Vec<Option<ExemplarSymbol>>,
}

impl<'a, SA, ISA, PSI> ExemplarDoc<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    fn from_documents(docs: &'a [&'a PsiDocument<'a, SA, ISA, PSI>]) -> Vec<Self> {
        docs.iter().copied().map(Self::new).collect()
    }

    fn new(document: &'a PsiDocument<'a, SA, ISA, PSI>) -> Self {
        let mut symbols = Vec::with_capacity(document.sigma.K().saturating_sub(1));
        let mut symbols_by_sigma = vec![None; document.sigma.K()];
        for symbol in 1..document.sigma.K() as u32 {
            let Some(text) = document.sigma.sigma_to_char(symbol) else {
                continue;
            };
            let Ok(range) = document.sigma.sa_range_for_sigma(symbol) else {
                continue;
            };
            if range.0 <= range.1 {
                let exemplar_symbol = ExemplarSymbol { text, range };
                symbols_by_sigma[symbol as usize] = Some(exemplar_symbol);
                symbols.push(exemplar_symbol);
            }
        }
        Self {
            document,
            symbols,
            symbols_by_sigma,
        }
    }

    fn constrain(
        &self,
        doc_state: &CorrelateDocState,
        symbol: ExemplarSymbol,
    ) -> Option<(usize, (usize, usize))> {
        let range = self
            .document
            .psi
            .constrain(&self.document.sigma, symbol.range, doc_state.range)
            .ok()?;
        if range.0 > range.1 {
            return None;
        }
        Some((range.1 - range.0 + 1, range))
    }
}

pub struct Exemplars<'a, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    docs: Vec<ExemplarDoc<'a, SA, ISA, PSI>>,
    heap: BinaryHeap<ExemplarState>,
    min_length: usize,
    sigma_ranges: Vec<(u32, (usize, usize))>,
}

impl<SA, ISA, PSI> Exemplars<'_, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    fn next_inner(&mut self) -> Option<ExemplarState> {
        while let Some(exemplar) = self.heap.pop() {
            if std::iter::zip(
                exemplar.stop.iter(),
                exemplar.needle[(exemplar.needle.len().saturating_sub(exemplar.stop.len()))..]
                    .iter(),
            )
            .all(|(x, y)| *x == *y)
                && exemplar.needle.len() >= self.min_length
            {
                return Some(exemplar);
            }
            if exemplar.docs.len() == 1 {
                let (idx, doc_state) = exemplar.docs.one().expect("len checked");
                let doc = &self.docs[idx];
                self.sigma_ranges.clear();
                if doc
                    .document
                    .psi
                    .predecessor_sigma_ranges(
                        &doc.document.sigma,
                        doc_state.range,
                        &mut self.sigma_ranges,
                    )
                    .unwrap_or(false)
                {
                    for (sigma, range) in self.sigma_ranges.iter().copied() {
                        let Some(symbol) =
                            doc.symbols_by_sigma.get(sigma as usize).copied().flatten()
                        else {
                            continue;
                        };
                        let numer = range.1 - range.0 + 1;
                        let mut needle = exemplar.needle.clone();
                        needle.push(symbol.text);
                        self.heap.push(ExemplarState {
                            stop: exemplar.stop.clone(),
                            needle,
                            cardinality: numer,
                            docs: ExemplarDocs::One(idx, CorrelateDocState { numer, range }),
                        });
                    }
                } else {
                    for symbol in doc.symbols.iter().copied() {
                        let Some((numer, range)) = doc.constrain(doc_state, symbol) else {
                            continue;
                        };
                        let mut needle = exemplar.needle.clone();
                        needle.push(symbol.text);
                        self.heap.push(ExemplarState {
                            stop: exemplar.stop.clone(),
                            needle,
                            cardinality: numer,
                            docs: ExemplarDocs::One(idx, CorrelateDocState { numer, range }),
                        });
                    }
                }
                continue;
            }
            let mut new_exemplars: HashMap<u32, ExemplarState> = HashMap::new();
            for (idx, doc_state) in exemplar.docs.iter() {
                let doc = &self.docs[idx];
                self.sigma_ranges.clear();
                if doc
                    .document
                    .psi
                    .predecessor_sigma_ranges(
                        &doc.document.sigma,
                        doc_state.range,
                        &mut self.sigma_ranges,
                    )
                    .unwrap_or(false)
                {
                    for (sigma, range) in self.sigma_ranges.iter().copied() {
                        let Some(symbol) =
                            doc.symbols_by_sigma.get(sigma as usize).copied().flatten()
                        else {
                            continue;
                        };
                        let numer = range.1 - range.0 + 1;
                        let new_exemplar = new_exemplars.entry(symbol.text).or_insert_with(|| {
                            let mut needle = exemplar.needle.clone();
                            needle.push(symbol.text);
                            ExemplarState {
                                stop: exemplar.stop.clone(),
                                needle,
                                cardinality: 0,
                                docs: ExemplarDocs::with_capacity(exemplar.docs.len()),
                            }
                        });
                        new_exemplar.cardinality = new_exemplar.cardinality.saturating_add(numer);
                        new_exemplar
                            .docs
                            .insert(idx, CorrelateDocState { numer, range });
                    }
                } else {
                    for symbol in doc.symbols.iter().copied() {
                        let Some((numer, range)) = doc.constrain(doc_state, symbol) else {
                            continue;
                        };
                        let new_exemplar = new_exemplars.entry(symbol.text).or_insert_with(|| {
                            let mut needle = exemplar.needle.clone();
                            needle.push(symbol.text);
                            ExemplarState {
                                stop: exemplar.stop.clone(),
                                needle,
                                cardinality: 0,
                                docs: ExemplarDocs::with_capacity(exemplar.docs.len()),
                            }
                        });
                        new_exemplar.cardinality = new_exemplar.cardinality.saturating_add(numer);
                        new_exemplar
                            .docs
                            .insert(idx, CorrelateDocState { numer, range });
                    }
                }
            }
            for (_, new_exemplar) in new_exemplars {
                if !new_exemplar.docs.is_empty() {
                    self.heap.push(new_exemplar);
                }
            }
        }
        None
    }
}

impl<SA, ISA, PSI> Iterator for Exemplars<'_, SA, ISA, PSI>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
{
    type Item = Exemplar;

    fn next(&mut self) -> Option<Self::Item> {
        self.next_inner().map(Exemplar::from)
    }
}

pub fn correlate<'a, SA, ISA, PSI, F>(
    docs: &'a [&'a PsiDocument<'a, SA, ISA, PSI>],
    boundaries: &[(u32, u32)],
    select: F,
) -> Correlate<'a, SA, ISA, PSI, F>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
    F: Fn(usize, RecordOffset) -> bool,
{
    let exemplars = exemplars(docs, boundaries);
    let correlates = BinaryHeap::new();
    Correlate {
        docs,
        exemplars,
        correlates,
        select,
    }
}

pub struct Correlate<'a, SA, ISA, PSI, F>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
    F: Fn(usize, RecordOffset) -> bool,
{
    docs: &'a [&'a PsiDocument<'a, SA, ISA, PSI>],
    exemplars: Exemplars<'a, SA, ISA, PSI>,
    correlates: BinaryHeap<ExemplarState>,
    select: F,
}

impl<SA, ISA, PSI, F> Iterator for Correlate<'_, SA, ISA, PSI, F>
where
    SA: sa::SuffixArray,
    ISA: isa::InverseSuffixArray,
    PSI: psi::Psi,
    F: Fn(usize, RecordOffset) -> bool,
{
    type Item = Exemplar;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(mut exemplar) = self.exemplars.next_inner() {
                let cardinality = exemplar.cardinality();
                exemplar.cardinality = 0;
                for (idx, doc) in exemplar.docs.iter_mut() {
                    doc.numer = 0;
                    for offset in doc.range.0..=doc.range.1 {
                        // TODO(rescrv): no unwrap
                        let offset = TextOffset(
                            self.docs[idx]
                                .sa
                                .lookup(&self.docs[idx].sigma, &self.docs[idx].psi, offset)
                                .unwrap(),
                        );
                        // TODO(rescrv): no unwrap
                        let offset = self.docs[idx].lookup(offset).unwrap();
                        if (self.select)(idx, offset) {
                            doc.numer += 1;
                        }
                    }
                    exemplar.cardinality = exemplar.cardinality.saturating_add(doc.numer);
                }
                self.correlates.push(exemplar);
                if let Some(correlate) = self.correlates.peek()
                    && correlate.cardinality() >= cardinality
                {
                    return self.correlates.pop().map(Exemplar::from);
                }
            } else if let Some(correlate) = self.correlates.pop() {
                return Some(Exemplar::from(correlate));
            } else {
                return None;
            }
        }
    }
}

////////////////////////////////////////////// inverse /////////////////////////////////////////////

fn inverse(x: &[usize]) -> Vec<usize> {
    let mut ix = vec![0usize; x.len()];
    for i in 0..x.len() {
        ix[x[i]] = i
    }
    ix
}

fn inverse_and_psi_u32(sa: &[u32]) -> (Vec<u32>, Vec<u32>) {
    let len = sa.len();
    let mut isa = vec![u32::MAX; len];
    let mut psi: Vec<std::mem::MaybeUninit<u32>> = Vec::with_capacity(len);
    // SAFETY(codex): MaybeUninit<u32> has no initialization invariant, and set_len only exposes
    // uninitialized slots that this function writes before converting to Vec<u32>. We never read a
    // slot until it has been written, which follows the Rustonomicon's MaybeUninit rules.
    unsafe {
        psi.set_len(len);
    }
    for (idx, pos) in sa.iter().copied().enumerate() {
        let idx = idx as u32;
        let pos = pos as usize;
        isa[pos] = idx;

        let prev_pos = if pos == 0 { len - 1 } else { pos - 1 };
        // SAFETY(codex): prev_pos is either pos - 1 or len - 1 and pos came from the suffix array,
        // so prev_pos < len == isa.len(); isa contains initialized u32 sentinels, making this
        // immutable read in-bounds and alias-safe under the Rustonomicon.
        let prev_isa = unsafe { *isa.get_unchecked(prev_pos) };
        if prev_isa != u32::MAX {
            // SAFETY(codex): prev_isa was read from isa after being initialized to a suffix-array
            // index, so prev_isa < len; each psi slot is written exactly when its predecessor is
            // discovered, and MaybeUninit::write avoids dropping uninitialized memory.
            unsafe {
                psi.get_unchecked_mut(prev_isa as usize).write(idx);
            }
        }

        let next_pos = if pos + 1 == len { 0 } else { pos + 1 };
        // SAFETY(codex): next_pos is either pos + 1 within len or 0, so it is in-bounds for isa;
        // isa is fully initialized with sentinel values before the loop and shared reads do not
        // violate Rustonomicon aliasing rules.
        let next_isa = unsafe { *isa.get_unchecked(next_pos) };
        if next_isa != u32::MAX {
            // SAFETY(codex): idx is the enumerate index converted to u32 with len < u32::MAX on
            // callers of this path, so idx < psi.len(); the slot is written with MaybeUninit::write
            // exactly when its successor is known, satisfying initialization requirements.
            unsafe {
                psi.get_unchecked_mut(idx as usize).write(next_isa);
            }
        }
    }
    // SAFETY(codex): valid suffix arrays are permutations, so the predecessor/successor loop writes
    // every psi slot exactly once before conversion; assume_init_vec_u32 only reinterprets fully
    // initialized u32 storage and preserves pointer, length, and capacity.
    (isa, unsafe { assume_init_vec_u32(psi) })
}

unsafe fn assume_init_vec_u32(mut values: Vec<std::mem::MaybeUninit<u32>>) -> Vec<u32> {
    let ptr = values.as_mut_ptr() as *mut u32;
    let len = values.len();
    let capacity = values.capacity();
    std::mem::forget(values);
    // SAFETY(codex): the caller guarantees every MaybeUninit<u32> element is initialized; u32 has
    // identical layout in MaybeUninit<u32>, and Vec::from_raw_parts receives the original pointer,
    // length, and capacity, satisfying the Rustonomicon ownership and allocation rules.
    unsafe { Vec::from_raw_parts(ptr, len, capacity) }
}

fn translate_text<Sym>(text: &[u32], sigma: &sigma::Sigma<'_>) -> Result<Vec<Sym>, Error>
where
    Sym: From<u8> + TryFrom<u32>,
{
    let mut s = Vec::with_capacity(text.len() + 1);
    for t in text.iter().copied() {
        let symbol = sigma.char_to_sigma(t).ok_or(Error::LogicError(
            "freshly constructed sigma cannot translate text",
        ))?;
        s.push(Sym::try_from(symbol).map_err(|_| Error::InvalidSigma)?);
    }
    s.push(Sym::from(0));
    Ok(s)
}

fn translate_text_u8(text: &[u32], sigma: &sigma::Sigma<'_>) -> Result<Vec<u8>, Error> {
    let mut s = Vec::with_capacity(text.len() + 1);
    if let Some(dense) = sigma.dense_char_to_sigma_table() {
        for t in text.iter().copied() {
            let symbol = dense.get(t as usize).copied().unwrap_or(0);
            if symbol == 0 || symbol > u8::MAX as u32 {
                return Err(Error::InvalidSigma);
            }
            s.push(symbol as u8);
        }
    } else {
        for t in text.iter().copied() {
            let symbol = sigma.char_to_sigma(t).ok_or(Error::LogicError(
                "freshly constructed sigma cannot translate text",
            ))?;
            s.push(u8::try_from(symbol).map_err(|_| Error::InvalidSigma)?);
        }
    }
    s.push(0);
    Ok(s)
}

fn translate_text_u16(text: &[u32], sigma: &sigma::Sigma<'_>) -> Result<Vec<u16>, Error> {
    let mut s = Vec::with_capacity(text.len() + 1);
    if let Some(dense) = sigma.dense_char_to_sigma_table() {
        for t in text.iter().copied() {
            let symbol = dense.get(t as usize).copied().unwrap_or(0);
            if symbol == 0 || symbol > u16::MAX as u32 {
                return Err(Error::InvalidSigma);
            }
            s.push(symbol as u16);
        }
    } else {
        for t in text.iter().copied() {
            let symbol = sigma.char_to_sigma(t).ok_or(Error::LogicError(
                "freshly constructed sigma cannot translate text",
            ))?;
            s.push(u16::try_from(symbol).map_err(|_| Error::InvalidSigma)?);
        }
    }
    s.push(0);
    Ok(s)
}

////////////////////////////////////// check_record_boundaries /////////////////////////////////////

fn check_record_boundaries(text: &[u32], record_boundaries: &[usize]) -> Result<(), Error> {
    if record_boundaries.is_empty() {
        return Err(Error::BadRecordOffset);
    }
    for (lhs, rhs) in std::iter::zip(record_boundaries.iter(), record_boundaries[1..].iter()) {
        if lhs >= rhs {
            return Err(Error::BadRecordOffset);
        }
    }
    if record_boundaries[0] != 0 {
        return Err(Error::BadRecordOffset);
    }
    if record_boundaries[record_boundaries.len() - 1] >= text.len() {
        return Err(Error::BadRecordOffset);
    }
    Ok(())
}

/////////////////////////////////////////////// utils //////////////////////////////////////////////

fn bits_required(max: u64) -> u8 {
    let max = std::cmp::max(max, 1);
    let max = max.next_power_of_two();
    let bits = max.ilog2() as u8;
    bits + 1
}

///////////////////////////////////////////// test_util ////////////////////////////////////////////

#[cfg(test)]
pub mod test_util {
    use crate::builder::Builder;
    use crate::sigma::Sigma;

    #[macro_export]
    macro_rules! assert_eq_with_ctx {
        (@inner [$($elems:tt)*] , $($rem:tt)*) => {
            format!("{} = {:?}; {}", stringify!($($elems)*), $($elems)*, assert_eq_with_ctx!(@inner [] $($rem)*))
        };

        (@inner [$($elems:tt)*] $e:tt $($rem:tt)*) => {
            assert_eq_with_ctx!(@inner [$($elems)* $e] $($rem)*)
        };

        (@inner [$($elems:tt)*]) => {
            format!("{} = {:?}", stringify!($($elems)*), $($elems)*)
        };

        ($lhs:expr, $rhs:expr, $($rem:expr),+) => {
            assert_eq!($lhs, $rhs, "{} == {}; {}", stringify!($lhs), stringify!($rhs), assert_eq_with_ctx!(@inner [] $($rem),*));
        };

        ($lhs:expr, $rhs:expr) => {
            assert_eq!($lhs, $rhs, "{} == {}", stringify!($lhs), stringify!($rhs));
        };
    }

    pub(crate) use assert_eq_with_ctx;

    // NOTE(rescrv):  Verbose, possibly duplicating information, so that there's a canonical place
    // to look for the examples of different representations.  Also serves as a human cross-check.
    // Static, possibly to a point of inconvenience, to make the most easy thing to do to create
    // top-level constants.  These are both done to serve human readers well and force human
    // writers to think about the test case and generate it by hand.
    #[allow(non_snake_case)]
    pub struct TestCase {
        pub text: &'static str,
        pub sigma2text: &'static [char],
        pub boundaries: &'static [usize],
        pub not_in_str: &'static [char],
        pub S: &'static [u32],
        pub SA: &'static [usize],
        pub ISA: &'static [usize],
        pub PSI: &'static [usize],
        pub bucket_starts: &'static [usize],
        pub bucket_limits: &'static [usize],
        pub deref_SA: &'static [u32],
        pub lstype: &'static str,
        pub lmspos: &'static str,
        pub searches: &'static [(&'static str, &'static [usize])],
        pub table: &'static str,
        #[allow(clippy::type_complexity)]
        pub constrain: &'static [((usize, usize), (usize, usize), (usize, usize))],
    }

    impl TestCase {
        pub fn sigma(&self) -> Vec<u8> {
            let mut buf = Vec::new();
            let mut builder = Builder::new(&mut buf);
            Sigma::construct(self.text.chars().map(|c| c as u32), &mut builder)
                .expect("test case should construct");
            drop(builder);
            buf
        }
    }

    pub const BANANA: &TestCase = &TestCase {
        text: "BANANA",
        sigma2text: &['A', 'B', 'N'],
        boundaries: &[1, 4, 5, 7],
        not_in_str: &['C', 'D', 'E'],
        S: &[2, 1, 3, 1, 3, 1, 0],
        SA: &[6, 5, 3, 1, 0, 4, 2],
        ISA: &[4, 3, 6, 2, 5, 1, 0],
        PSI: &[4, 0, 5, 6, 3, 1, 2],
        bucket_starts: &[0, 1, 4, 5],
        bucket_limits: &[1, 4, 5, 7],
        deref_SA: &[0, 1, 1, 1, 2, 3, 3],
        lstype: "LSLSLLS",
        lmspos: " * *  *",
        searches: &[("AN", &[1, 3]), ("NA", &[2, 4])],
        table: "
    A      B   N
[]  [0]    []  [] 
[]  []     []  [1]
[]  []     [3] [2]
[4] []     []  [] 
[]  [5, 6] []  [] 
",
        constrain: &[
            // Backwards search of AN
            ((5, 6), (0, 7), (5, 6)),
            ((1, 3), (5, 6), (2, 3)),
            // Backwards search of BANA
            ((1, 3), (0, 7), (1, 3)),
            ((5, 6), (1, 3), (5, 6)),
            ((1, 3), (5, 6), (2, 3)),
            ((4, 4), (2, 3), (4, 4)),
        ],
    };

    pub const MISSISSIPPI: &TestCase = &TestCase {
        text: "MISSISSIPPI",
        sigma2text: &['I', 'M', 'P', 'S'],
        boundaries: &[1, 5, 6, 8, 12],
        not_in_str: &['A', 'B', 'N'],
        S: &[2, 1, 4, 4, 1, 4, 4, 1, 3, 3, 1, 0],
        SA: &[11, 10, 7, 4, 1, 0, 9, 8, 6, 3, 5, 2],
        ISA: &[5, 4, 11, 9, 3, 10, 8, 2, 7, 6, 1, 0],
        PSI: &[5, 0, 7, 10, 11, 4, 1, 6, 2, 3, 8, 9],
        bucket_starts: &[0, 1, 5, 6, 8],
        bucket_limits: &[1, 5, 6, 8, 12],
        deref_SA: &[0, 1, 1, 1, 1, 2, 3, 3, 4, 4, 4, 4],
        lstype: "LSLLSLLSLLLS",
        lmspos: " *  *  *   *",
        searches: &[("ISS", &[1, 4])],
        table: "
    I        M   P   S
[]  [0]      []  []  []
[]  []       []  [1] []
[]  []       []  []  [2]
[]  []       [4] []  [3]
[5] []       []  []  []
[]  []       []  [6] []
[]  [7]      []  []  []
[]  []       []  []  [8, 9]
[]  [10, 11] []  []  []
",
        constrain: &[
            ((8, 11), (8, 11), (10, 11)),
            ((8, 11), (1, 4), (8, 9)),
            ((1, 4), (8, 11), (3, 4)),
            ((1, 4), (0, 0), (1, 1)),
            ((1, 4), (6, 7), (2, 2)),
            ((5, 5), (1, 4), (5, 5)),
            ((6, 7), (6, 7), (7, 7)),
        ],
    };

    pub const MISSISSIPPI_BANANA: &TestCase = &TestCase {
        text: "MISSISSIPPIBANANA",
        sigma2text: &['A', 'B', 'I', 'M', 'N', 'P', 'S'],
        boundaries: &[1, 4, 5, 9, 10, 12, 14, 18],
        not_in_str: &['C', 'D', 'E'],
        S: &[4, 3, 7, 7, 3, 7, 7, 3, 6, 6, 3, 2, 1, 5, 1, 5, 1, 0],
        SA: &[17, 16, 14, 12, 11, 10, 7, 4, 1, 0, 15, 13, 9, 8, 6, 3, 5, 2],
        ISA: &[9, 8, 17, 15, 7, 16, 14, 6, 13, 12, 5, 4, 3, 11, 2, 10, 1, 0],
        PSI: &[9, 0, 10, 11, 3, 4, 13, 16, 17, 8, 1, 2, 5, 12, 6, 7, 14, 15],
        bucket_starts: &[0, 1, 4, 5, 9, 10, 12, 14],
        bucket_limits: &[1, 4, 5, 9, 10, 12, 14, 18],
        deref_SA: &[0, 1, 1, 1, 2, 3, 3, 3, 3, 4, 5, 5, 6, 6, 7, 7, 7, 7],
        lstype: "LSLLSLLSLLLLSLSLLS",
        lmspos: " *  *  *    * *  *",
        searches: &[("ISS", &[1, 4])],
        table: "
    A        B   I        M   N   P    S
[]  [0]      []  []       []  []  []   []
[]  []       []  []       []  [1] []   []
[]  []       [3] []       []  [2] []   []
[]  []       []  [4]      []  []  []   []
[]  []       []  []       []  []  [5]  []
[]  []       []  []       []  []  []   [6]
[]  []       []  []       [8] []  []   [7]
[9] []       []  []       []  []  []   []
[]  [10, 11] []  []       []  []  []   []
[]  []       []  []       []  []  [12] []
[]  []       []  [13]     []  []  []   []
[]  []       []  []       []  []  []   [14, 15]
[]  []       []  [16, 17] []  []  []   []
",
        constrain: &[],
    };

    pub const MIISSISSISSIPPI: &TestCase = &TestCase {
        text: "MIISSISSISSIPPI",
        sigma2text: &['I', 'M', 'P', 'S'],
        boundaries: &[1, 7, 8, 10, 16],
        not_in_str: &['A', 'B', 'N'],
        S: &[2, 1, 1, 4, 4, 1, 4, 4, 1, 4, 4, 1, 3, 3, 1, 0],
        SA: &[15, 14, 1, 11, 8, 5, 2, 0, 13, 12, 10, 7, 4, 9, 6, 3],
        ISA: &[7, 2, 6, 15, 12, 5, 14, 11, 4, 13, 10, 3, 9, 8, 1, 0],
        PSI: &[7, 0, 6, 9, 13, 14, 15, 2, 1, 8, 3, 4, 5, 10, 11, 12],
        bucket_starts: &[0, 1, 7, 8, 10],
        bucket_limits: &[1, 7, 8, 10, 16],
        deref_SA: &[0, 1, 1, 1, 1, 1, 1, 2, 3, 3, 4, 4, 4, 4, 4, 4],
        lstype: "LSSLLSLLSLLSLLLS",
        lmspos: " *   *  *  *   *",
        searches: &[("ISS", &[2, 5, 8])],
        table: "
    I            M   P   S
[]  [0]          []  []  []
[]  []           []  [1] []
[]  []           [2] []  []
[]  []           []  []  [3]
[]  [6]          []  []  [4, 5]
[7] []           []  []  []
[]  []           []  [8] []
[]  [9]          []  []  []
[]  []           []  []  [10, 11, 12]
[]  [13, 14, 15] []  []  []
",
        constrain: &[],
    };

    // Thist test case was found in early prototype of sais.  It has survived a language rewrite
    // and really only exists because it makes me laugh.  I'd love to replace it with a minimal
    // test case demonstrating all of the following:
    // - identical S+LS blocks ordered in SA the way they are ordered in the input
    // - identical S+LS blocks reversed in SA from the way they are ordered in the input
    // - S+LS block that whose S+L prefix is shared by multiple S+L strings, also forward/reverse
    pub const MUTANT_BANANA: &TestCase = &TestCase {
        text: "ANBNANANCANBNCNANANCNA",
        sigma2text: &['A', 'B', 'C', 'N'],
        boundaries: &[1, 8, 10, 13, 23],
        not_in_str: &['D', 'E', 'F'],
        S: &[
            1, 4, 2, 4, 1, 4, 1, 4, 3, 1, 4, 2, 4, 3, 4, 1, 4, 1, 4, 3, 4, 1, 0,
        ],
        SA: &[
            22, 21, 4, 15, 0, 9, 6, 17, 2, 11, 8, 19, 13, 20, 3, 14, 5, 16, 1, 10, 7, 18, 12,
        ],
        ISA: &[
            4, 18, 8, 14, 2, 16, 6, 20, 10, 5, 19, 9, 22, 12, 15, 3, 17, 7, 21, 11, 13, 1, 0,
        ],
        PSI: &[
            4, 0, 16, 17, 18, 19, 20, 21, 14, 22, 5, 13, 15, 1, 2, 3, 6, 7, 8, 9, 10, 11, 12,
        ],
        bucket_starts: &[0, 1, 8, 10, 13],
        bucket_limits: &[1, 8, 10, 13, 23],
        deref_SA: &[
            0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
        ],
        lstype: "SLSLSLSLLSLSLSLSLSLSLLS",
        lmspos: "  * * *  * * * * * *  *",
        searches: &[("AN", &[0, 4, 6, 9, 15, 17])],
        table: "
    A        B    C        N
[]  [0]      []   []       []
[]  []       []   []       [1]
[4] []       []   [5]      [2, 3, 6, 7]
[]  []       []   []       [8, 9]
[]  []       []   []       [10]
[]  []       []   []       [11, 12]
[]  [16, 17] [14] [13, 15] []
[]  [18, 19] []   []       []
[]  [20, 21] [22] []       []
",
        constrain: &[],
    };

    pub const SAIS_EXAMPLE: &TestCase = &TestCase {
        text: "baabababbab",
        sigma2text: &['a', 'b'],
        boundaries: &[1, 6, 12],
        not_in_str: &['c', 'd', 'e'],
        S: &[2, 1, 1, 2, 1, 2, 1, 2, 2, 1, 2, 0],
        SA: &[11, 1, 9, 2, 4, 6, 10, 0, 8, 3, 5, 7],
        ISA: &[7, 1, 3, 9, 4, 10, 5, 11, 8, 2, 6, 0],
        PSI: &[7, 3, 6, 9, 10, 11, 0, 1, 2, 4, 5, 8],
        bucket_starts: &[0, 1, 6],
        bucket_limits: &[1, 6, 12],
        deref_SA: &[0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
        lstype: "LSSLSLSLLSLS",
        lmspos: " *  * *  * *",
        searches: &[("b", &[0, 3, 5, 7, 8, 10])],
        table: "
    a       b
[]  []      [0]
[]  []      [1]
[]  [3]     [2, 4, 5]
[]  [6]     []
[7] [9, 10] [8]
[]  [11]    []
",
        constrain: &[],
    };

    pub const BAD_CASE_1: &TestCase = &TestCase {
        text: "1111111111112111",
        sigma2text: &['1', '2'],
        boundaries: &[1, 16, 17],
        not_in_str: &['3'],
        S: &[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 0],
        SA: &[16, 15, 14, 13, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
        ISA: &[4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 3, 2, 1, 0],
        PSI: &[4, 0, 1, 2, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 3],
        bucket_starts: &[0, 1, 16],
        bucket_limits: &[1, 16, 17],
        deref_SA: &[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2],
        lstype: "SSSSSSSSSSSSLLLLS",
        lmspos: "                *",
        searches: &[],
        table: "
    1                                      2
[]  [0]                                    []
[]  [1]                                    []
[4] [2, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14] [3]
[]  [15]                                   []
[]  [16]                                   []
",
        constrain: &[((1, 15), (16, 16), (15, 15))],
    };

    pub const BAD_CASE_2: &TestCase = &TestCase {
        text: "2121221211111132",
        sigma2text: &['1', '2', '3'],
        boundaries: &[1, 10, 16, 17],
        not_in_str: &['4'],
        S: &[2, 1, 2, 1, 2, 2, 1, 2, 1, 1, 1, 1, 1, 1, 3, 2, 0],
        SA: &[16, 8, 9, 10, 11, 12, 6, 1, 3, 13, 15, 7, 5, 0, 2, 4, 14],
        ISA: &[13, 7, 14, 8, 15, 12, 6, 11, 1, 2, 3, 4, 5, 9, 16, 10, 0],
        PSI: &[13, 2, 3, 4, 5, 9, 11, 14, 15, 16, 0, 1, 6, 7, 8, 12, 10],
        bucket_starts: &[0, 1, 10, 16],
        bucket_limits: &[1, 10, 16, 17],
        deref_SA: &[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3],
        lstype: "LSLSLLSLSSSSSSLLS",
        lmspos: " * *  * *       *",
        searches: &[],
        table: "
     1            2         3
[]   []           [0]       []
[]   [2, 3, 4, 5] [1]       []
[]   []           [6, 7, 8] []
[]   [9]          []        []
[]   []           []        [10]
[13] [11, 14]     [12]      []
[]   [15]         []        []
[]   [16]         []        []
",
        constrain: &[((16, 16), (1, 9), (16, 15))],
    };

    pub const BAD_CASE_3: &TestCase = &TestCase {
        text: "2121221212121222",
        sigma2text: &['1', '2'],
        boundaries: &[1, 7, 17],
        not_in_str: &['3'],
        S: &[2, 1, 2, 1, 2, 2, 1, 2, 1, 2, 1, 2, 1, 2, 2, 2, 0],
        SA: &[16, 6, 8, 1, 10, 3, 12, 15, 5, 7, 0, 9, 2, 11, 14, 4, 13],
        ISA: &[10, 3, 12, 5, 15, 8, 1, 9, 2, 11, 4, 13, 6, 16, 14, 7, 0],
        PSI: &[10, 9, 11, 12, 13, 15, 16, 0, 1, 2, 3, 4, 5, 6, 7, 8, 14],
        bucket_starts: &[0, 1, 7],
        bucket_limits: &[1, 7, 17],
        deref_SA: &[0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
        lstype: "LSLSLLSLSLSLSLLLS",
        lmspos: " * *  * * * *   *",
        searches: &[],
        table: "
     1               2
[]   []              [0]
[]   []              [1, 2, 3, 4, 5, 6]
[]   []              [7]
[10] [9, 11, 12, 13] [8]
[]   [15, 16]        [14]
",
        constrain: &[((1, 6), (1, 6), (1, 0))],
    };

    pub const BAD_CASE_4: &TestCase = &TestCase {
        text: "12224122231222312222122221222212221122211222112221",
        sigma2text: &['1', '2', '3', '4'],
        boundaries: &[1, 15, 48, 50, 51],
        not_in_str: &['5'],
        S: &[1, 2, 2, 2, 4, 1, 2, 2, 2, 3, 1, 2, 2, 2, 3, 1, 2, 2, 2, 2, 1, 2, 2, 2, 2, 1, 2, 2, 2, 2, 1, 2, 2, 2, 1, 1, 2, 2, 2, 1, 1, 2, 2, 2, 1, 1, 2, 2, 2, 1, 0],
        SA: &[50, 49, 44, 39, 34, 45, 40, 35, 30, 25, 20, 15, 10, 5, 0, 48, 43, 38, 33, 29, 24, 19, 47, 42, 37, 32, 28, 23, 18, 46, 41, 36, 31, 27, 22, 17, 26, 21, 16, 11, 6, 1, 12, 7, 2, 13, 8, 3, 14, 9, 4],
        ISA: &[14, 41, 44, 47, 50, 13, 40, 43, 46, 49, 12, 39, 42, 45, 48, 11, 38, 35, 28, 21, 10, 37, 34, 27, 20, 9, 36, 33, 26, 19, 8, 32, 25, 18, 4, 7, 31, 24, 17, 3, 6, 30, 23, 16, 2, 5, 29, 22, 15, 1, 0],
        PSI: &[14, 0, 5, 6, 7, 29, 30, 31, 32, 36, 37, 38, 39, 40, 41, 1, 2, 3, 4, 8, 9, 10, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 33, 34, 35, 42, 43, 44, 45, 46, 47, 48, 49, 50, 11, 12, 13],
        bucket_starts: &[0, 1, 15, 48, 50],
        bucket_limits: &[1, 15, 48, 50, 51],
        deref_SA: &[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 4],
        lstype: "SSSSLSSSSLSSSSLSLLLLSLLLLSLLLLSLLLSSLLLSSLLLSSLLLLS",
        lmspos: "     *    *    *    *    *    *   *    *    *     *",
        searches: &[],
        table: "
     1                                        2                                                    3        4
[]   [0]                                      []                                                   []       []
[]   []                                       [1]                                                  []       []
[]   []                                       [2, 3, 4]                                            []       []
[14] [5, 6, 7]                                [8, 9, 10]                                           [11, 12] [13]
[]   []                                       [15, 16, 17, 18, 19, 20, 21]                         []       []
[]   [29, 30, 31, 32, 36, 37, 38, 39, 40, 41] [22, 23, 24, 25, 26, 27, 28, 33, 34, 35, 42, 43, 44] []       []
[]   []                                       [45, 46]                                             []       []
[]   []                                       [47]                                                 []       []
[]   []                                       [48, 49]                                             []       []
[]   []                                       [50]                                                 []       []
",
        constrain: &[
            ((15, 47), (48, 49), (45, 46)),
            ((15, 47), (45, 46), (42, 43)),
            ((15, 47), (42, 43), (39, 40)),
            ((15, 47), (50, 50), (47, 47)),
            ((15, 47), (47, 47), (44, 44)),
        ],
    };

    #[macro_export]
    macro_rules! test_cases_for {
        ($name:ident, $check:path) => {
            mod $name {
                #[test]
                fn banana() {
                    $check($crate::test_util::BANANA);
                }

                #[test]
                fn mississippi() {
                    $check($crate::test_util::MISSISSIPPI);
                }

                #[test]
                fn mississippi_banana() {
                    $check($crate::test_util::MISSISSIPPI_BANANA);
                }

                #[test]
                fn miissississippi() {
                    $check($crate::test_util::MIISSISSISSIPPI);
                }

                #[test]
                fn mutant_banana() {
                    $check($crate::test_util::MUTANT_BANANA);
                }

                #[test]
                fn sais_example() {
                    $check($crate::test_util::SAIS_EXAMPLE);
                }

                #[test]
                fn bad_case1() {
                    $check($crate::test_util::BAD_CASE_1);
                }

                #[test]
                fn bad_case2() {
                    $check($crate::test_util::BAD_CASE_2);
                }

                #[test]
                fn bad_case3() {
                    $check($crate::test_util::BAD_CASE_3);
                }

                #[test]
                fn bad_case4() {
                    $check($crate::test_util::BAD_CASE_4);
                }
            }
        };
    }

    pub(crate) use test_cases_for;

    #[macro_export]
    macro_rules! search_cases_for {
        ($name:ident, $check:path) => {
            mod $name {
                use super::*;

                fn build(t: &TestCase) -> Vec<u8> {
                    let text: Vec<u32> = t.text.chars().map(|c| c as u32).collect();
                    let record_boundaries = vec![0usize];
                    let mut buf = Vec::new();
                    let mut builder = Builder::new(&mut buf);
                    <$check>::construct(text, record_boundaries, &mut builder)
                        .expect("test case should index");
                    drop(builder);
                    buf
                }

                fn check_len(t: &TestCase) {
                    let doc = build(t);
                    let (doc, _) =
                        <$check>::unpack(&doc).expect("freshly created document should unpack");
                    assert_eq!(t.text.chars().count(), doc.len());
                }

                test_cases_for! {len, super::check_len}

                fn check_retrieve(t: &TestCase) {
                    let text: Vec<u32> = t.text.chars().map(|c| c as u32).collect();
                    let doc = build(t);
                    let (doc, _) =
                        <$check>::unpack(&doc).expect("freshly created document should unpack");
                    assert_eq!(
                        text,
                        doc.retrieve(RecordOffset(0))
                            .expect("retrieve should succeed")
                    );
                }

                test_cases_for! {retrieve, super::check_retrieve}

                fn check_search(t: &TestCase) {
                    let doc = build(t);
                    let (doc, _) =
                        <$check>::unpack(&doc).expect("freshly created document should unpack");
                    for (needle, offsets) in t.searches.iter() {
                        let expected: Vec<TextOffset> =
                            offsets.iter().map(|x| TextOffset(*x)).collect();
                        let needle_u32: Vec<u32> = needle.chars().map(|c| c as u32).collect();
                        let mut returned: Vec<TextOffset> = doc
                            .search(&needle_u32)
                            .expect("search should succeed")
                            .collect();
                        returned.sort();
                        assert_eq_with_ctx!(expected, returned, needle);
                    }
                    // TODO(rescrv): Check empty search.
                }

                test_cases_for! {search, super::check_search}

                fn check_count(t: &TestCase) {
                    let doc = build(t);
                    let (doc, _) =
                        <$check>::unpack(&doc).expect("freshly created document should unpack");
                    for (needle, offsets) in t.searches.iter() {
                        let needle_u32: Vec<u32> = needle.chars().map(|c| c as u32).collect();
                        let count: usize = doc.count(&needle_u32).expect("count should succeed");
                        assert_eq_with_ctx!(offsets.len(), count, needle);
                    }
                }

                test_cases_for! {count, super::check_count}
            }
        };
    }

    pub(crate) use search_cases_for;
}

#[cfg(test)]
mod tests {
    use crate::isa::InverseSuffixArray;
    use crate::sa::SuffixArray;

    use super::psi::Psi;
    use super::test_util::*;
    use super::*;

    #[test]
    #[allow(clippy::assertions_on_constants)]
    fn u32_as_usize() {
        assert!(u32::BITS <= usize::BITS);
    }

    #[test]
    #[allow(clippy::assertions_on_constants)]
    fn usize_as_u64() {
        assert!(usize::BITS <= u64::BITS);
    }

    #[test]
    fn inverse() {
        let x = &[8, 6, 9, 5, 0, 3, 1, 2, 7, 4];
        let ix = &[4, 6, 7, 5, 9, 3, 1, 8, 0, 2];
        let returned: &[usize] = &super::inverse(x);
        assert_eq!(ix, returned);
        let returned: &[usize] = &super::inverse(returned);
        assert_eq!(x, returned);
    }

    fn check_string(t: &TestCase) {
        let sigma = t.sigma();
        let sigma = sigma::Sigma::unpack(&sigma).expect("test should unpack").0;
        let mut s: Vec<u32> = t
            .text
            .chars()
            .map(|x| sigma.char_to_sigma(x as u32).unwrap())
            .collect();
        s.push(0);
        assert_eq!(t.S, s);
    }

    test_cases_for! {string, super::check_string}

    fn check_isa(t: &TestCase) {
        let returned = super::inverse(t.SA);
        assert_eq_with_ctx!(t.ISA, returned);
    }

    test_cases_for! {isa, super::check_isa}

    fn check_unpack_psi_document(t: &TestCase) {
        let text: Vec<u32> = t.text.chars().map(|c| c as u32).collect();
        let record_boundaries = vec![0usize];
        let mut buf = Vec::new();
        let mut builder = Builder::new(&mut buf);
        type Document<'a> = super::PsiDocument<
            'a,
            super::sa::ReferenceSuffixArray,
            super::isa::ReferenceInverseSuffixArray,
            super::psi::ReferencePsi,
        >;
        Document::construct(text, record_boundaries, &mut builder).expect("test case should index");
        drop(builder);
        let doc = Document::unpack(&buf).expect("document should unpack").0;

        for (idx, sa) in t.SA.iter().enumerate() {
            assert_eq_with_ctx!(*sa, doc.sa.lookup(&doc.sigma, &doc.psi, idx).unwrap(), idx);
        }

        for (idx, isa) in t.ISA.iter().enumerate() {
            assert_eq_with_ctx!(*isa, doc.isa.lookup(idx).unwrap(), idx);
        }

        assert_eq!(t.PSI.len(), doc.psi.len());
        for (idx, psi) in t.PSI.iter().enumerate() {
            assert_eq_with_ctx!(*psi, doc.psi.lookup(&doc.sigma, idx).unwrap(), idx);
        }
    }

    test_cases_for! {unpack_psi_document, super::check_unpack_psi_document}

    search_cases_for! {reference, crate::ReferenceDocument}
    search_cases_for! {psi_with_all_reference, crate::PsiDocument::<crate::sa::ReferenceSuffixArray, crate::isa::ReferenceInverseSuffixArray, psi::ReferencePsi>}
    search_cases_for! {psi_with_wavelet_psi, crate::PsiDocument::<crate::sa::ReferenceSuffixArray, crate::isa::ReferenceInverseSuffixArray, psi::wavelet_tree::WaveletTreePsi<wavelet_tree::ReferenceWaveletTree>>}
    search_cases_for! {compressed_document, crate::CompressedDocument}
}
