dino_quant/

lib.rs

use std::f32::consts::PI;
use std::{io::Read, path::Path};

const BASE_A: u8 = 0;
const BASE_C: u8 = 1;
const BASE_G: u8 = 2;
const BASE_T: u8 = 3;
const CONTIG_SEPARATOR: u8 = b'N';

#[derive(Clone, Copy, Debug)]
pub struct SplitMix64 {
    state: u64,
}

impl SplitMix64 {
    pub const fn new(seed: u64) -> Self {
        Self { state: seed }
    }

    pub fn next_u64(&mut self) -> u64 {
        self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15);
        mix_u64(self.state)
    }
}

#[derive(Clone, Copy, Debug)]
pub struct QuantizerConfig {
    pub bits: u8,
    pub clip_sigma: f32,
    pub rotation_seed: u64,
    pub qjl_seed: u64,
    pub use_qjl_residual: bool,
}

impl Default for QuantizerConfig {
    fn default() -> Self {
        Self {
            bits: 4,
            clip_sigma: 4.0,
            rotation_seed: 0x51d0_51d0_51d0_51d0,
            qjl_seed: 0x71b0_7171_b071_7171,
            use_qjl_residual: true,
        }
    }
}

#[derive(Clone, Debug)]
pub struct QjlResidual {
    signs: Vec<u64>,
    dim: usize,
    norm: f32,
    seed: u64,
}

#[derive(Clone, Debug)]
struct PackedCodes {
    data: Vec<u8>,
    len: usize,
    bits: u8,
}

#[derive(Clone, Debug)]
pub struct QuantizedVector {
    dim: usize,
    bits: u8,
    clip: f32,
    rotation_seed: u64,
    codes: PackedCodes,
    qjl_residual: Option<QjlResidual>,
}

#[derive(Clone, Debug)]
pub struct PreparedQuantizedQuery {
    dim: usize,
    rotation_seed: u64,
    qjl_seed: Option<u64>,
    rotated: Vec<f32>,
    rotated_sum: f32,
    scalar_4bit_lookup: Vec<[f32; 256]>,
    rotated_qjl: Option<Vec<f32>>,
}

#[derive(Clone, Debug)]
pub struct QuantizedVectorSnapshot {
    pub dim: usize,
    pub bits: u8,
    pub clip: f32,
    pub rotation_seed: u64,
    pub codes_data: Vec<u8>,
    pub codes_len: usize,
    pub codes_bits: u8,
    pub qjl_residual: Option<QjlResidualSnapshot>,
}

#[derive(Clone, Debug)]
pub struct QjlResidualSnapshot {
    pub signs: Vec<u64>,
    pub dim: usize,
    pub norm: f32,
    pub seed: u64,
}

#[derive(Clone, Copy, Debug)]
pub struct ReconstructionMetrics {
    pub mse: f32,
    pub cosine: f32,
    pub dot_error: f32,
}

#[derive(Clone, Debug)]
pub struct SequenceRecord {
    pub name: String,
    pub bases: Vec<u8>,
}

#[derive(Clone, Copy, Debug)]
pub struct SequenceRecordRef<'a> {
    pub name: &'a [u8],
    pub bases: &'a [u8],
}

#[derive(Clone, Copy, Debug)]
pub struct ReferenceIndexConfig {
    pub k: usize,
    pub dim: usize,
    pub window_len: usize,
    pub stride: usize,
    pub quantizer: QuantizerConfig,
}

#[derive(Clone, Debug)]
pub struct SearchHit {
    pub target_name: String,
    pub target_start: usize,
    pub target_end: usize,
    pub start: usize,
    pub end: usize,
    pub score: f32,
}

#[derive(Clone, Debug)]
pub struct ReferenceWindowIndex {
    config: ReferenceIndexConfig,
    windows: Vec<ReferenceWindow>,
}

#[derive(Clone, Debug)]
struct ReferenceWindow {
    target_name: String,
    target_start: usize,
    target_end: usize,
    start: usize,
    end: usize,
    sketch: QuantizedVector,
}

impl QuantizerConfig {
    pub fn encode(&self, vector: &[f32]) -> Result<QuantizedVector, String> {
        self.validate(vector.len())?;

        let dim = vector.len();
        let clip = self.clip_sigma / (dim as f32).sqrt();
        let rotated = rotate(vector, self.rotation_seed)?;
        let codes = encode_scalar_codes(&rotated, self.bits, clip)?;

        let mut quantized = QuantizedVector {
            dim,
            bits: self.bits,
            clip,
            rotation_seed: self.rotation_seed,
            codes,
            qjl_residual: None,
        };

        if self.use_qjl_residual {
            let base = quantized.decode()?;
            let residual = subtract(vector, &base)?;
            quantized.qjl_residual = QjlResidual::encode(&residual, self.qjl_seed)?;
        }

        Ok(quantized)
    }

    fn validate(&self, dim: usize) -> Result<(), String> {
        if dim == 0 {
            return Err("vector dimension must be non-zero".to_owned());
        }
        if !dim.is_power_of_two() {
            return Err(
                "vector dimension must be a power of two for the Hadamard rotation".to_owned(),
            );
        }
        if !(1..=8).contains(&self.bits) {
            return Err("quantizer bits must be in 1..=8".to_owned());
        }
        if !self.clip_sigma.is_finite() || self.clip_sigma <= 0.0 {
            return Err("clip_sigma must be finite and positive".to_owned());
        }
        Ok(())
    }
}

impl QuantizedVector {
    pub fn prepare_approximate_query(
        query: &[f32],
        rotation_seed: u64,
        qjl_seed: Option<u64>,
    ) -> Result<PreparedQuantizedQuery, String> {
        let rotated = rotate(query, rotation_seed)?;
        let rotated_sum = rotated.iter().sum();
        let scalar_4bit_lookup = build_4bit_lookup(&rotated);
        let rotated_qjl = qjl_seed.map(|seed| rotate(query, seed)).transpose()?;
        Ok(PreparedQuantizedQuery {
            dim: query.len(),
            rotation_seed,
            qjl_seed,
            rotated,
            rotated_sum,
            scalar_4bit_lookup,
            rotated_qjl,
        })
    }

    pub fn snapshot(&self) -> QuantizedVectorSnapshot {
        QuantizedVectorSnapshot {
            dim: self.dim,
            bits: self.bits,
            clip: self.clip,
            rotation_seed: self.rotation_seed,
            codes_data: self.codes.data.clone(),
            codes_len: self.codes.len,
            codes_bits: self.codes.bits,
            qjl_residual: self
                .qjl_residual
                .as_ref()
                .map(|residual| QjlResidualSnapshot {
                    signs: residual.signs.clone(),
                    dim: residual.dim,
                    norm: residual.norm,
                    seed: residual.seed,
                }),
        }
    }

    pub fn from_snapshot(snapshot: QuantizedVectorSnapshot) -> Result<Self, String> {
        if snapshot.dim == 0 || !snapshot.dim.is_power_of_two() {
            return Err(
                "quantized vector snapshot dimension must be a non-zero power of two".to_owned(),
            );
        }
        if snapshot.codes_len != snapshot.dim {
            return Err("quantized vector snapshot code length must match dimension".to_owned());
        }
        if !(1..=8).contains(&snapshot.bits) || !(1..=8).contains(&snapshot.codes_bits) {
            return Err("quantized vector snapshot bits must be in 1..=8".to_owned());
        }
        if snapshot.bits != snapshot.codes_bits {
            return Err("quantized vector snapshot code bits mismatch".to_owned());
        }
        if !snapshot.clip.is_finite() || snapshot.clip <= 0.0 {
            return Err("quantized vector snapshot clip must be finite and positive".to_owned());
        }
        let expected_code_bytes =
            (snapshot.codes_len * usize::from(snapshot.codes_bits)).div_ceil(8);
        if snapshot.codes_data.len() != expected_code_bytes {
            return Err("quantized vector snapshot code byte length mismatch".to_owned());
        }
        let qjl_residual = snapshot
            .qjl_residual
            .map(|residual| {
                if residual.dim != snapshot.dim {
                    return Err(
                        "QJL snapshot dimension must match quantized vector dimension".to_owned(),
                    );
                }
                if residual.signs.len() != residual.dim.div_ceil(64) {
                    return Err("QJL snapshot sign word length mismatch".to_owned());
                }
                if !residual.norm.is_finite() || residual.norm <= 0.0 {
                    return Err("QJL snapshot norm must be finite and positive".to_owned());
                }
                Ok(QjlResidual {
                    signs: residual.signs,
                    dim: residual.dim,
                    norm: residual.norm,
                    seed: residual.seed,
                })
            })
            .transpose()?;
        Ok(Self {
            dim: snapshot.dim,
            bits: snapshot.bits,
            clip: snapshot.clip,
            rotation_seed: snapshot.rotation_seed,
            codes: PackedCodes {
                data: snapshot.codes_data,
                len: snapshot.codes_len,
                bits: snapshot.codes_bits,
            },
            qjl_residual,
        })
    }

    pub fn decode(&self) -> Result<Vec<f32>, String> {
        if self.codes.len() != self.dim {
            return Err("quantized code length does not match dimension".to_owned());
        }

        let rotated = decode_scalar_codes(&self.codes, self.bits, self.clip)?;
        let mut decoded = inverse_rotate(&rotated, self.rotation_seed)?;
        if let Some(residual) = &self.qjl_residual {
            let correction = residual.decode()?;
            if correction.len() != decoded.len() {
                return Err("QJL correction length does not match decoded vector".to_owned());
            }
            for (dst, corr) in decoded.iter_mut().zip(correction) {
                *dst += corr;
            }
        }
        Ok(decoded)
    }

    pub fn compressed_bits(&self) -> usize {
        let scalar_bits = self.codes.byte_len() * 8;
        let qjl_bits = self
            .qjl_residual
            .as_ref()
            .map_or(0, |residual| residual.dim + 32);
        scalar_bits + qjl_bits
    }

    pub fn compressed_bytes(&self) -> usize {
        self.compressed_bits().div_ceil(8)
    }

    pub fn scalar_dot_rotated_query(&self, rotated_query: &[f32]) -> Result<f32, String> {
        if rotated_query.len() != self.dim {
            return Err("rotated query dimension does not match quantized vector".to_owned());
        }
        let rotated_query_sum = rotated_query.iter().sum();
        self.scalar_dot_rotated_query_with_sum(rotated_query, rotated_query_sum)
    }

    fn scalar_dot_rotated_query_with_sum(
        &self,
        rotated_query: &[f32],
        rotated_query_sum: f32,
    ) -> Result<f32, String> {
        if rotated_query.len() != self.dim {
            return Err("rotated query dimension does not match quantized vector".to_owned());
        }
        self.codes
            .decoded_dot(rotated_query, rotated_query_sum, self.clip)
    }

    fn approximate_dot_4bit_lookup(
        &self,
        lookup: &[[f32; 256]],
        rotated_query_sum: f32,
        rotated_qjl_query: Option<&[f32]>,
    ) -> Result<f32, String> {
        if self.bits != 4 {
            return Err("4-bit lookup scoring requires a 4-bit quantized vector".to_owned());
        }
        let scale = 2.0 * self.clip / 15.0;
        let scalar_score = (-self.clip * rotated_query_sum)
            + scale * self.codes.weighted_code_sum_4_lookup(lookup)?;
        Ok(scalar_score + self.residual_dot_rotated_query(rotated_qjl_query)?)
    }

    fn residual_dot_rotated_query(&self, rotated_qjl_query: Option<&[f32]>) -> Result<f32, String> {
        match (&self.qjl_residual, rotated_qjl_query) {
            (Some(residual), Some(query)) => residual.dot_rotated_query(query),
            (Some(_), None) => Err("QJL residual scoring requires a rotated QJL query".to_owned()),
            (None, _) => Ok(0.0),
        }
    }

    pub fn scalar_dot_query(&self, query: &[f32]) -> Result<f32, String> {
        let rotated_query = rotate(query, self.rotation_seed)?;
        self.scalar_dot_rotated_query(&rotated_query)
    }

    pub fn approximate_dot_query(&self, query: &[f32]) -> Result<f32, String> {
        let rotated_query = rotate(query, self.rotation_seed)?;
        let rotated_query_sum = rotated_query.iter().sum();
        let mut score =
            self.scalar_dot_rotated_query_with_sum(&rotated_query, rotated_query_sum)?;
        if let Some(residual) = &self.qjl_residual {
            let rotated_residual_query = rotate(query, residual.seed)?;
            score += residual.dot_rotated_query(&rotated_residual_query)?;
        }
        Ok(score)
    }

    pub fn approximate_dot_prepared_query(
        &self,
        query: &PreparedQuantizedQuery,
    ) -> Result<f32, String> {
        if query.dim != self.dim {
            return Err("prepared query dimension does not match quantized vector".to_owned());
        }
        if query.rotation_seed != self.rotation_seed {
            return Err("prepared query rotation seed does not match quantized vector".to_owned());
        }
        let rotated_qjl = if let Some(residual) = &self.qjl_residual {
            if query.qjl_seed != Some(residual.seed) {
                return Err("prepared query QJL seed does not match residual".to_owned());
            }
            Some(
                query
                    .rotated_qjl
                    .as_deref()
                    .ok_or_else(|| "prepared query is missing QJL rotation".to_owned())?,
            )
        } else {
            None
        };
        if self.bits == 4 {
            return self.approximate_dot_4bit_lookup(
                &query.scalar_4bit_lookup,
                query.rotated_sum,
                rotated_qjl,
            );
        }
        let mut score =
            self.scalar_dot_rotated_query_with_sum(&query.rotated, query.rotated_sum)?;
        if let Some(residual) = &self.qjl_residual {
            let Some(rotated_qjl) = rotated_qjl else {
                return Err("prepared query is missing QJL rotation".to_owned());
            };
            score += residual.dot_rotated_query(rotated_qjl)?;
        }
        Ok(score)
    }
}

impl PackedCodes {
    fn encode(values: &[u8], bits: u8) -> Result<Self, String> {
        if !(1..=8).contains(&bits) {
            return Err("packed code bits must be in 1..=8".to_owned());
        }
        let mask = (1_u16 << bits) - 1;
        let mut data = vec![0_u8; (values.len() * usize::from(bits)).div_ceil(8)];
        for (idx, code) in values.iter().enumerate() {
            if u16::from(*code) > mask {
                return Err("scalar code exceeds bit width".to_owned());
            }
            let bit_pos = idx * usize::from(bits);
            let byte_idx = bit_pos / 8;
            let bit_offset = bit_pos % 8;
            let shifted = u16::from(*code) << bit_offset;
            data[byte_idx] |= shifted as u8;
            if bit_offset + usize::from(bits) > 8 {
                data[byte_idx + 1] |= (shifted >> 8) as u8;
            }
        }

        Ok(Self {
            data,
            len: values.len(),
            bits,
        })
    }

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

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

    fn get(&self, idx: usize) -> Result<u8, String> {
        if idx >= self.len {
            return Err("packed code index out of bounds".to_owned());
        }
        let bit_pos = idx * usize::from(self.bits);
        let byte_idx = bit_pos / 8;
        let bit_offset = bit_pos % 8;
        let mut value = u16::from(self.data[byte_idx] >> bit_offset);
        if bit_offset + usize::from(self.bits) > 8 {
            value |= u16::from(self.data[byte_idx + 1]) << (8 - bit_offset);
        }
        Ok((value & ((1_u16 << self.bits) - 1)) as u8)
    }

    fn decoded_dot(
        &self,
        rotated_query: &[f32],
        rotated_query_sum: f32,
        clip: f32,
    ) -> Result<f32, String> {
        if rotated_query.len() != self.len {
            return Err("rotated query length does not match packed code length".to_owned());
        }
        if !clip.is_finite() || clip <= 0.0 {
            return Err("clip must be finite and positive".to_owned());
        }

        let levels = (1_u16 << self.bits) - 1;
        let scale = 2.0 * clip / levels as f32;
        let offset_sum = -clip * rotated_query_sum;
        Ok(offset_sum + scale * self.weighted_code_sum(rotated_query)?)
    }

    fn weighted_code_sum(&self, rotated_query: &[f32]) -> Result<f32, String> {
        match self.bits {
            2 => Ok(self.weighted_code_sum_2(rotated_query)),
            4 => Ok(self.weighted_code_sum_4(rotated_query)),
            8 => Ok(self.weighted_code_sum_8(rotated_query)),
            _ => self.weighted_code_sum_generic(rotated_query),
        }
    }

    fn weighted_code_sum_2(&self, rotated_query: &[f32]) -> f32 {
        let mut sum = 0.0_f32;
        for (byte_idx, byte) in self.data.iter().enumerate() {
            let idx = byte_idx * 4;
            if idx >= self.len {
                break;
            }
            sum += f32::from(byte & 0b0000_0011) * rotated_query[idx];
            if idx + 1 < self.len {
                sum += f32::from((byte >> 2) & 0b0000_0011) * rotated_query[idx + 1];
            }
            if idx + 2 < self.len {
                sum += f32::from((byte >> 4) & 0b0000_0011) * rotated_query[idx + 2];
            }
            if idx + 3 < self.len {
                sum += f32::from(byte >> 6) * rotated_query[idx + 3];
            }
        }
        sum
    }

    fn weighted_code_sum_4(&self, rotated_query: &[f32]) -> f32 {
        let mut sum = 0.0_f32;
        for (byte_idx, byte) in self.data.iter().enumerate() {
            let idx = byte_idx * 2;
            if idx >= self.len {
                break;
            }
            sum += f32::from(byte & 0x0f) * rotated_query[idx];
            if idx + 1 < self.len {
                sum += f32::from(byte >> 4) * rotated_query[idx + 1];
            }
        }
        sum
    }

    fn weighted_code_sum_4_lookup(&self, lookup: &[[f32; 256]]) -> Result<f32, String> {
        if lookup.len() != self.data.len() {
            return Err("4-bit lookup table length does not match packed code bytes".to_owned());
        }
        let mut sum = 0.0_f32;
        for (idx, byte) in self.data.iter().enumerate() {
            sum += lookup[idx][usize::from(*byte)];
        }
        Ok(sum)
    }

    fn weighted_code_sum_8(&self, rotated_query: &[f32]) -> f32 {
        let mut sum = 0.0_f32;
        for (code, query_value) in self.data.iter().zip(rotated_query) {
            sum += f32::from(*code) * query_value;
        }
        sum
    }

    fn weighted_code_sum_generic(&self, rotated_query: &[f32]) -> Result<f32, String> {
        let mut sum = 0.0_f32;
        let mask = (1_u32 << self.bits) - 1;
        let mut byte_idx = 0;
        let mut bit_buffer = 0_u32;
        let mut bits_in_buffer = 0_u8;
        for query_value in rotated_query {
            while bits_in_buffer < self.bits {
                if byte_idx >= self.data.len() {
                    return Err("packed code buffer ended early".to_owned());
                }
                bit_buffer |= u32::from(self.data[byte_idx]) << bits_in_buffer;
                bits_in_buffer += 8;
                byte_idx += 1;
            }

            let code = bit_buffer & mask;
            bit_buffer >>= self.bits;
            bits_in_buffer -= self.bits;
            sum += code as f32 * query_value;
        }
        Ok(sum)
    }

    fn decode_all(&self, clip: f32) -> Result<Vec<f32>, String> {
        if !clip.is_finite() || clip <= 0.0 {
            return Err("clip must be finite and positive".to_owned());
        }
        let levels = (1_u16 << self.bits) - 1;
        let scale = 2.0 * clip / levels as f32;
        let offset = -clip;
        let mut decoded = Vec::with_capacity(self.len);
        for idx in 0..self.len {
            decoded.push(offset + f32::from(self.get(idx)?) * scale);
        }
        Ok(decoded)
    }
}

impl ReferenceIndexConfig {
    pub fn validate(&self) -> Result<(), String> {
        if self.window_len == 0 {
            return Err("window_len must be non-zero".to_owned());
        }
        if self.stride == 0 {
            return Err("stride must be non-zero".to_owned());
        }
        if self.k == 0 || self.k > 31 {
            return Err("k must be in 1..=31".to_owned());
        }
        self.quantizer.validate(self.dim)
    }
}

impl ReferenceWindowIndex {
    pub fn build(reference: &[u8], config: ReferenceIndexConfig) -> Result<Self, String> {
        config.validate()?;
        if reference.len() < config.window_len {
            return Err("reference is shorter than the configured window length".to_owned());
        }

        let window_count = ((reference.len() - config.window_len) / config.stride) + 1;
        let mut windows = Vec::with_capacity(window_count);
        for start in (0..=reference.len() - config.window_len).step_by(config.stride) {
            let end = start + config.window_len;
            let sketch = dna_kmer_sketch(&reference[start..end], config.k, config.dim)?;
            let sketch = config.quantizer.encode(&sketch)?;
            windows.push(ReferenceWindow {
                target_name: "reference".to_owned(),
                target_start: start,
                target_end: end,
                start,
                end,
                sketch,
            });
        }

        Ok(Self { config, windows })
    }

    pub fn build_records(
        records: &[SequenceRecord],
        config: ReferenceIndexConfig,
    ) -> Result<Self, String> {
        config.validate()?;
        if records.is_empty() {
            return Err("at least one sequence record is required".to_owned());
        }

        let mut windows = Vec::new();
        let mut linear_offset = 0_usize;
        for (idx, record) in records.iter().enumerate() {
            if record.bases.len() >= config.window_len {
                let window_count = ((record.bases.len() - config.window_len) / config.stride) + 1;
                windows.reserve(window_count);
                for target_start in
                    (0..=record.bases.len() - config.window_len).step_by(config.stride)
                {
                    let target_end = target_start + config.window_len;
                    let sketch = dna_kmer_sketch(
                        &record.bases[target_start..target_end],
                        config.k,
                        config.dim,
                    )?;
                    let sketch = config.quantizer.encode(&sketch)?;
                    let start = linear_offset + target_start;
                    let end = linear_offset + target_end;
                    windows.push(ReferenceWindow {
                        target_name: record.name.clone(),
                        target_start,
                        target_end,
                        start,
                        end,
                        sketch,
                    });
                }
            }
            linear_offset += record.bases.len();
            if idx + 1 < records.len() {
                linear_offset += 1;
            }
        }

        if windows.is_empty() {
            return Err("no reference record is long enough for the configured window".to_owned());
        }

        Ok(Self { config, windows })
    }

    pub fn search_sequence(&self, query: &[u8], top_k: usize) -> Result<Vec<SearchHit>, String> {
        let sketch = dna_kmer_sketch(query, self.config.k, self.config.dim)?;
        self.search_sketch(&sketch, top_k)
    }

    pub fn search_sketch(&self, query: &[f32], top_k: usize) -> Result<Vec<SearchHit>, String> {
        if top_k == 0 {
            return Ok(Vec::new());
        }
        if query.len() != self.config.dim {
            return Err("query sketch dimension does not match index dimension".to_owned());
        }

        let rotated_query = rotate(query, self.config.quantizer.rotation_seed)?;
        let rotated_query_sum = rotated_query.iter().sum();
        let rotated_qjl_query = if self.config.quantizer.use_qjl_residual {
            Some(rotate(query, self.config.quantizer.qjl_seed)?)
        } else {
            None
        };
        if self.config.quantizer.bits == 4 {
            return self.search_rotated_4bit_lookup(
                &rotated_query,
                rotated_query_sum,
                rotated_qjl_query.as_deref(),
                top_k,
            );
        }

        let mut top = Vec::with_capacity(top_k.min(self.windows.len()));
        for window in &self.windows {
            let scalar_score = window
                .sketch
                .scalar_dot_rotated_query_with_sum(&rotated_query, rotated_query_sum)?;
            let score = scalar_score
                + window
                    .sketch
                    .residual_dot_rotated_query(rotated_qjl_query.as_deref())?;
            push_top_hit(
                &mut top,
                top_k,
                SearchHit {
                    target_name: window.target_name.clone(),
                    target_start: window.target_start,
                    target_end: window.target_end,
                    start: window.start,
                    end: window.end,
                    score,
                },
            );
        }
        top.sort_by(|left, right| right.score.total_cmp(&left.score));
        Ok(top)
    }

    fn search_rotated_4bit_lookup(
        &self,
        rotated_query: &[f32],
        rotated_query_sum: f32,
        rotated_qjl_query: Option<&[f32]>,
        top_k: usize,
    ) -> Result<Vec<SearchHit>, String> {
        let lookup = build_4bit_lookup(rotated_query);
        let mut top = Vec::with_capacity(top_k.min(self.windows.len()));
        for window in &self.windows {
            let score = window.sketch.approximate_dot_4bit_lookup(
                &lookup,
                rotated_query_sum,
                rotated_qjl_query,
            )?;
            push_top_hit(
                &mut top,
                top_k,
                SearchHit {
                    target_name: window.target_name.clone(),
                    target_start: window.target_start,
                    target_end: window.target_end,
                    start: window.start,
                    end: window.end,
                    score,
                },
            );
        }
        top.sort_by(|left, right| right.score.total_cmp(&left.score));
        Ok(top)
    }

    pub fn window_count(&self) -> usize {
        self.windows.len()
    }

    pub fn compressed_bytes(&self) -> usize {
        self.windows
            .iter()
            .map(|window| window.sketch.compressed_bytes())
            .sum()
    }

    pub fn raw_vector_bytes(&self) -> usize {
        self.windows.len() * self.config.dim * std::mem::size_of::<f32>()
    }

    pub fn compression_ratio(&self) -> f32 {
        if self.compressed_bytes() == 0 {
            return 0.0;
        }
        self.raw_vector_bytes() as f32 / self.compressed_bytes() as f32
    }

    pub fn config(&self) -> ReferenceIndexConfig {
        self.config
    }
}

impl QjlResidual {
    fn encode(residual: &[f32], seed: u64) -> Result<Option<Self>, String> {
        let norm = l2_norm(residual);
        if norm <= f32::EPSILON {
            return Ok(None);
        }

        let projected = rotate(residual, seed)?;
        let mut signs = vec![0_u64; projected.len().div_ceil(64)];
        for (idx, value) in projected.iter().enumerate() {
            if *value >= 0.0 {
                signs[idx / 64] |= 1_u64 << (idx % 64);
            }
        }

        Ok(Some(Self {
            signs,
            dim: projected.len(),
            norm,
            seed,
        }))
    }

    fn decode(&self) -> Result<Vec<f32>, String> {
        if self.dim == 0 || !self.dim.is_power_of_two() {
            return Err("QJL residual dimension must be a non-zero power of two".to_owned());
        }
        let scale = (PI / 2.0).sqrt() * self.norm / self.dim as f32;
        let mut projected = vec![0.0_f32; self.dim];
        for (idx, slot) in projected.iter_mut().enumerate() {
            let bit = (self.signs[idx / 64] >> (idx % 64)) & 1;
            *slot = if bit == 1 { scale } else { -scale };
        }
        inverse_rotate(&projected, self.seed)
    }

    fn dot_rotated_query(&self, rotated_query: &[f32]) -> Result<f32, String> {
        if rotated_query.len() != self.dim {
            return Err("QJL rotated query dimension does not match residual dimension".to_owned());
        }
        let scale = (PI / 2.0).sqrt() * self.norm / self.dim as f32;
        let mut sum = 0.0_f32;
        for (idx, query_value) in rotated_query.iter().enumerate() {
            let bit = (self.signs[idx / 64] >> (idx % 64)) & 1;
            let sign = if bit == 1 { 1.0 } else { -1.0 };
            sum += sign * query_value;
        }
        Ok(scale * sum)
    }
}

pub fn rotate(vector: &[f32], seed: u64) -> Result<Vec<f32>, String> {
    validate_hadamard_dim(vector.len())?;
    let mut rotated = Vec::with_capacity(vector.len());
    for (idx, value) in vector.iter().enumerate() {
        rotated.push(*value * coordinate_sign(seed, idx));
    }
    hadamard_in_place(&mut rotated);
    let scale = 1.0 / (vector.len() as f32).sqrt();
    for value in &mut rotated {
        *value *= scale;
    }
    Ok(rotated)
}

pub fn inverse_rotate(rotated: &[f32], seed: u64) -> Result<Vec<f32>, String> {
    validate_hadamard_dim(rotated.len())?;
    let mut vector = rotated.to_vec();
    hadamard_in_place(&mut vector);
    let scale = 1.0 / (rotated.len() as f32).sqrt();
    for (idx, value) in vector.iter_mut().enumerate() {
        *value *= scale * coordinate_sign(seed, idx);
    }
    Ok(vector)
}

pub fn dna_kmer_sketch(seq: &[u8], k: usize, dim: usize) -> Result<Vec<f32>, String> {
    if k == 0 || k > 31 {
        return Err("k must be in 1..=31".to_owned());
    }
    if dim == 0 {
        return Err("sketch dimension must be non-zero".to_owned());
    }
    if seq.len() < k {
        return Ok(vec![0.0; dim]);
    }

    let mut sketch = vec![0.0_f32; dim];
    for window in seq.windows(k) {
        if let Some(code) = canonical_kmer_code(window) {
            let hash = mix_u64(code ^ ((k as u64) << 56));
            let bucket = (hash as usize) % dim;
            let sign = if (hash >> 63) == 0 { 1.0 } else { -1.0 };
            sketch[bucket] += sign;
        }
    }
    l2_normalize(&mut sketch);
    Ok(sketch)
}

pub fn protein_kmer_sketch(seq: &[u8], k: usize, dim: usize) -> Result<Vec<f32>, String> {
    if k == 0 || k > 16 {
        return Err("protein k must be in 1..=16".to_owned());
    }
    if dim == 0 {
        return Err("sketch dimension must be non-zero".to_owned());
    }
    if seq.len() < k {
        return Ok(vec![0.0; dim]);
    }

    let mut sketch = vec![0.0_f32; dim];
    for window in seq.windows(k) {
        if let Some(code) = protein_kmer_code(window) {
            let hash = mix_u64(code ^ ((k as u64) << 56));
            let bucket = (hash as usize) % dim;
            let sign = if (hash >> 63) == 0 { 1.0 } else { -1.0 };
            sketch[bucket] += sign;
        }
    }
    l2_normalize(&mut sketch);
    Ok(sketch)
}

#[cfg(test)]
fn parse_fasta_bytes(input: &[u8]) -> Result<Vec<SequenceRecord>, String> {
    let mut records = Vec::new();
    dino_seq::visit_fasta_bytes(input, |record: dino_seq::FastaVisitRecord<'_>| {
        let mut bases = Vec::with_capacity(record.seq().len());
        append_sequence_line(record.seq(), &mut bases).map_err(dino_seq_format_error)?;
        records.push(SequenceRecord {
            name: parse_record_name(record.name_without_gt()).map_err(dino_seq_format_error)?,
            bases,
        });
        Ok(())
    })
    .map_err(|err| format!("failed to parse FASTA bytes with dino-seq: {err}"))?;
    if records.is_empty() {
        return Err("FASTA input did not contain any records".to_owned());
    }
    Ok(records)
}

pub fn read_fasta_file(path: impl AsRef<Path>) -> Result<Vec<SequenceRecord>, String> {
    let path = path.as_ref();
    let mut reader = dino_seq::open_fasta_for_reference(path).map_err(|err| {
        format!(
            "failed to open FASTA with dino-seq {}: {err}",
            path.display()
        )
    })?;
    let mut records = Vec::new();
    reader
        .visit_records(|record| {
            let mut bases = Vec::with_capacity(record.seq().len());
            append_sequence_line(record.seq(), &mut bases).map_err(dino_seq_format_error)?;
            records.push(SequenceRecord {
                name: parse_record_name(record.name_without_gt()).map_err(dino_seq_format_error)?,
                bases,
            });
            Ok(())
        })
        .map_err(|err| {
            format!(
                "failed to parse FASTA with dino-seq {}: {err}",
                path.display()
            )
        })?;
    if records.is_empty() {
        return Err("FASTA input did not contain any records".to_owned());
    }
    Ok(records)
}

pub fn read_protein_fasta_file(path: impl AsRef<Path>) -> Result<Vec<SequenceRecord>, String> {
    let path = path.as_ref();
    let mut reader = dino_seq::open_fasta(path).map_err(|err| {
        format!(
            "failed to open protein FASTA with dino-seq {}: {err}",
            path.display()
        )
    })?;
    let mut records = Vec::new();
    reader
        .visit_records(|record| {
            let mut bases = Vec::with_capacity(record.seq().len());
            append_protein_line(record.seq(), &mut bases).map_err(dino_seq_format_error)?;
            records.push(SequenceRecord {
                name: parse_record_name(record.name_without_gt()).map_err(dino_seq_format_error)?,
                bases,
            });
            Ok(())
        })
        .map_err(|err| {
            format!(
                "failed to parse protein FASTA with dino-seq {}: {err}",
                path.display()
            )
        })?;
    if records.is_empty() {
        return Err("protein FASTA input did not contain any records".to_owned());
    }
    Ok(records)
}

#[cfg(test)]
fn parse_fastq_bytes(input: &[u8]) -> Result<Vec<SequenceRecord>, String> {
    let mut records = Vec::new();
    dino_seq::visit_fastq_bytes(
        input,
        dino_seq::FastqConfig::default(),
        |record: dino_seq::FastqVisitRecord<'_>| {
            let header = record.name();
            let name = header.strip_prefix(b"@").unwrap_or(header);
            let mut bases = Vec::with_capacity(record.seq().len());
            append_sequence_line(record.seq(), &mut bases).map_err(dino_seq_format_error)?;
            records.push(SequenceRecord {
                name: parse_record_name(name).map_err(dino_seq_format_error)?,
                bases,
            });
            Ok(())
        },
    )
    .map_err(|err| format!("failed to parse FASTQ bytes with dino-seq: {err}"))?;
    if records.is_empty() {
        return Err("FASTQ input did not contain any records".to_owned());
    }
    Ok(records)
}

pub fn visit_fastq_slices_file(
    path: impl AsRef<Path>,
    visitor: impl FnMut(SequenceRecordRef<'_>) -> Result<(), String>,
) -> Result<(), String> {
    visit_fastq_slices_file_limit(path, None, visitor)
}

pub fn visit_fastq_slices_file_limit(
    path: impl AsRef<Path>,
    max_records: Option<usize>,
    visitor: impl FnMut(SequenceRecordRef<'_>) -> Result<(), String>,
) -> Result<(), String> {
    let path = path.as_ref();
    let mut reader = dino_seq::open_fastq(path).map_err(|err| {
        format!(
            "failed to open FASTQ with dino-seq {}: {err}",
            path.display()
        )
    })?;
    visit_fastq_slices_with_reader(&mut reader, max_records, visitor).map_err(|err| {
        format!(
            "failed to parse FASTQ with dino-seq {}: {err}",
            path.display()
        )
    })
}

fn visit_fastq_slices_with_reader<R: Read>(
    reader: &mut dino_seq::FastqReader<R>,
    max_records: Option<usize>,
    mut visitor: impl FnMut(SequenceRecordRef<'_>) -> Result<(), String>,
) -> Result<(), String> {
    let chunk_bytes = std::env::var("DINO_QUANT_FASTQ_CHUNK_BYTES")
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .filter(|&value| value > 0)
        .unwrap_or(4 * 1024 * 1024);
    let chunk_config = dino_seq::FastqChunkConfig::new(chunk_bytes).min_records(1);
    let mut records = 0_usize;
    while max_records.is_none_or(|limit| records < limit)
        && reader
            .next_chunk_with_sink(chunk_config, &mut |record: dino_seq::FastqVisitRecord<
                '_,
            >| {
                if max_records.is_some_and(|limit| records >= limit) {
                    return Ok(());
                }
                let header = record.name();
                let name = header.strip_prefix(b"@").unwrap_or(header);
                let name = parse_record_name_bytes(name).map_err(dino_seq_format_error)?;
                validate_sequence_bases(record.seq()).map_err(dino_seq_format_error)?;
                visitor(SequenceRecordRef {
                    name,
                    bases: record.seq(),
                })
                .map_err(dino_seq_format_error)?;
                records += 1;
                Ok(())
            })
            .map_err(|err| err.to_string())?
            .is_some()
    {}
    if records == 0 {
        return Err("FASTQ input did not contain any records".to_owned());
    }
    Ok(())
}

pub fn concatenate_records(records: &[SequenceRecord]) -> Result<Vec<u8>, String> {
    if records.is_empty() {
        return Err("at least one sequence record is required".to_owned());
    }
    let total_bases = records
        .iter()
        .map(|record| record.bases.len())
        .sum::<usize>();
    let mut concatenated = Vec::with_capacity(total_bases + records.len().saturating_sub(1));
    for (idx, record) in records.iter().enumerate() {
        if idx > 0 {
            concatenated.push(CONTIG_SEPARATOR);
        }
        concatenated.extend_from_slice(&record.bases);
    }
    Ok(concatenated)
}

pub fn reconstruction_metrics(
    original: &[f32],
    decoded: &[f32],
) -> Result<ReconstructionMetrics, String> {
    Ok(ReconstructionMetrics {
        mse: mse(original, decoded)?,
        cosine: cosine_similarity(original, decoded)?,
        dot_error: (dot(original, original)? - dot(original, decoded)?).abs(),
    })
}

pub fn mutate_dna(seq: &[u8], every: usize) -> Vec<u8> {
    if every == 0 {
        return seq.to_vec();
    }
    let mut mutated = seq.to_vec();
    for idx in (every - 1..mutated.len()).step_by(every) {
        mutated[idx] = match mutated[idx].to_ascii_uppercase() {
            b'A' => b'C',
            b'C' => b'G',
            b'G' => b'T',
            b'T' => b'A',
            other => other,
        };
    }
    mutated
}

pub fn synthetic_dna(len: usize, seed: u64) -> Vec<u8> {
    let mut rng = SplitMix64::new(seed);
    let mut seq = Vec::with_capacity(len);
    for _ in 0..len {
        let base = match rng.next_u64() & 3 {
            0 => b'A',
            1 => b'C',
            2 => b'G',
            _ => b'T',
        };
        seq.push(base);
    }
    seq
}

pub fn intervals_overlap(
    left_start: usize,
    left_end: usize,
    right_start: usize,
    right_end: usize,
) -> bool {
    left_start < right_end && right_start < left_end
}

pub fn dot(left: &[f32], right: &[f32]) -> Result<f32, String> {
    if left.len() != right.len() {
        return Err("vectors must have equal length".to_owned());
    }
    Ok(left.iter().zip(right).map(|(a, b)| a * b).sum())
}

pub fn mse(left: &[f32], right: &[f32]) -> Result<f32, String> {
    if left.len() != right.len() {
        return Err("vectors must have equal length".to_owned());
    }
    if left.is_empty() {
        return Err("vectors must be non-empty".to_owned());
    }
    let sum: f32 = left
        .iter()
        .zip(right)
        .map(|(a, b)| {
            let delta = a - b;
            delta * delta
        })
        .sum();
    Ok(sum / left.len() as f32)
}

pub fn cosine_similarity(left: &[f32], right: &[f32]) -> Result<f32, String> {
    let denom = l2_norm(left) * l2_norm(right);
    if denom <= f32::EPSILON {
        return Err("cosine similarity is undefined for zero vectors".to_owned());
    }
    Ok(dot(left, right)? / denom)
}

pub fn l2_normalize(vector: &mut [f32]) {
    let norm = l2_norm(vector);
    if norm <= f32::EPSILON {
        return;
    }
    for value in vector {
        *value /= norm;
    }
}

fn encode_scalar_codes(rotated: &[f32], bits: u8, clip: f32) -> Result<PackedCodes, String> {
    let levels = (1_u16 << bits) - 1;
    let inv_width = 1.0 / (2.0 * clip);
    let mut codes = Vec::with_capacity(rotated.len());
    for value in rotated {
        let clipped = value.clamp(-clip, clip);
        let normalized = (clipped + clip) * inv_width;
        codes.push((normalized * levels as f32).round() as u8);
    }
    PackedCodes::encode(&codes, bits)
}

fn decode_scalar_codes(codes: &PackedCodes, bits: u8, clip: f32) -> Result<Vec<f32>, String> {
    if !(1..=8).contains(&bits) {
        return Err("quantizer bits must be in 1..=8".to_owned());
    }
    if !clip.is_finite() || clip <= 0.0 {
        return Err("clip must be finite and positive".to_owned());
    }
    if codes.bits != bits {
        return Err("packed code bit width does not match quantizer bit width".to_owned());
    }
    codes.decode_all(clip)
}

fn subtract(left: &[f32], right: &[f32]) -> Result<Vec<f32>, String> {
    if left.len() != right.len() {
        return Err("vectors must have equal length".to_owned());
    }
    Ok(left.iter().zip(right).map(|(a, b)| a - b).collect())
}

fn push_top_hit(top: &mut Vec<SearchHit>, top_k: usize, hit: SearchHit) {
    if top.len() < top_k {
        top.push(hit);
        return;
    }

    let mut worst_idx = 0;
    let mut worst_score = top[0].score;
    for (idx, current) in top.iter().enumerate().skip(1) {
        if current.score < worst_score {
            worst_idx = idx;
            worst_score = current.score;
        }
    }

    if hit.score > worst_score {
        top[worst_idx] = hit;
    }
}

fn build_4bit_lookup(rotated_query: &[f32]) -> Vec<[f32; 256]> {
    let mut lookup = Vec::with_capacity(rotated_query.len().div_ceil(2));
    for chunk in rotated_query.chunks(2) {
        let first = chunk[0];
        let second = if chunk.len() == 2 { chunk[1] } else { 0.0 };
        let mut table = [0.0_f32; 256];
        for byte in 0_u16..=255 {
            let low = f32::from((byte & 0x0f) as u8);
            let high = f32::from((byte >> 4) as u8);
            table[usize::from(byte)] = low * first + high * second;
        }
        lookup.push(table);
    }
    lookup
}

fn trim_ascii(mut bytes: &[u8]) -> &[u8] {
    while matches!(bytes.first(), Some(b' ' | b'\t' | b'\r' | b'\n')) {
        bytes = &bytes[1..];
    }
    while matches!(bytes.last(), Some(b' ' | b'\t' | b'\r' | b'\n')) {
        bytes = &bytes[..bytes.len() - 1];
    }
    bytes
}

fn parse_record_name(header: &[u8]) -> Result<String, String> {
    let name = parse_record_name_bytes(header)?;
    String::from_utf8(name.to_vec()).map_err(|_| "sequence record name must be UTF-8".to_owned())
}

fn parse_record_name_bytes(header: &[u8]) -> Result<&[u8], String> {
    let trimmed = trim_ascii(header);
    if trimmed.is_empty() {
        return Err("sequence record header must contain a name".to_owned());
    }
    trimmed
        .split(|byte| matches!(*byte, b' ' | b'\t'))
        .next()
        .filter(|name| !name.is_empty())
        .ok_or_else(|| "sequence record header must contain a name".to_owned())
}

fn append_sequence_line(line: &[u8], dst: &mut Vec<u8>) -> Result<(), String> {
    for base in line {
        match base.to_ascii_uppercase() {
            b'A' | b'C' | b'G' | b'T' | b'N' => dst.push(base.to_ascii_uppercase()),
            b' ' | b'\t' | b'\r' => {}
            other => {
                return Err(format!(
                    "unsupported sequence character '{}' in input",
                    char::from(other)
                ));
            }
        }
    }
    Ok(())
}

fn append_protein_line(line: &[u8], dst: &mut Vec<u8>) -> Result<(), String> {
    for residue in line {
        let upper = residue.to_ascii_uppercase();
        match upper {
            b'A'..=b'Z' | b'*' | b'-' => dst.push(upper),
            b' ' | b'\t' | b'\r' => {}
            other => {
                return Err(format!(
                    "unsupported protein sequence character '{}' in input",
                    char::from(other)
                ));
            }
        }
    }
    Ok(())
}

fn validate_sequence_bases(line: &[u8]) -> Result<(), String> {
    for base in line {
        match base.to_ascii_uppercase() {
            b'A' | b'C' | b'G' | b'T' | b'N' => {}
            other => {
                return Err(format!(
                    "unsupported sequence character '{}' in input",
                    char::from(other)
                ));
            }
        }
    }
    Ok(())
}

fn protein_kmer_code(window: &[u8]) -> Option<u64> {
    let mut code = 0_u64;
    for residue in window {
        code = code.checked_mul(23)?;
        code = code.checked_add(u64::from(protein_code(*residue)?))?;
    }
    Some(code)
}

fn protein_code(residue: u8) -> Option<u8> {
    match residue.to_ascii_uppercase() {
        b'A' => Some(0),
        b'C' => Some(1),
        b'D' => Some(2),
        b'E' => Some(3),
        b'F' => Some(4),
        b'G' => Some(5),
        b'H' => Some(6),
        b'I' => Some(7),
        b'K' => Some(8),
        b'L' => Some(9),
        b'M' => Some(10),
        b'N' => Some(11),
        b'P' => Some(12),
        b'Q' => Some(13),
        b'R' => Some(14),
        b'S' => Some(15),
        b'T' => Some(16),
        b'V' => Some(17),
        b'W' => Some(18),
        b'Y' => Some(19),
        b'B' | b'J' | b'O' | b'U' | b'X' | b'Z' | b'*' | b'-' => None,
        _ => None,
    }
}

fn dino_seq_format_error(message: String) -> dino_seq::FastqError {
    dino_seq::FastqError::Format(message)
}

fn l2_norm(vector: &[f32]) -> f32 {
    vector.iter().map(|value| value * value).sum::<f32>().sqrt()
}

fn validate_hadamard_dim(dim: usize) -> Result<(), String> {
    if dim == 0 {
        return Err("Hadamard dimension must be non-zero".to_owned());
    }
    if !dim.is_power_of_two() {
        return Err("Hadamard dimension must be a power of two".to_owned());
    }
    Ok(())
}

fn hadamard_in_place(values: &mut [f32]) {
    let mut stride = 1;
    while stride < values.len() {
        let step = stride * 2;
        for start in (0..values.len()).step_by(step) {
            for idx in start..start + stride {
                let left = values[idx];
                let right = values[idx + stride];
                values[idx] = left + right;
                values[idx + stride] = left - right;
            }
        }
        stride = step;
    }
}

fn coordinate_sign(seed: u64, idx: usize) -> f32 {
    if mix_u64(seed ^ idx as u64) & 1 == 0 {
        1.0
    } else {
        -1.0
    }
}

fn canonical_kmer_code(window: &[u8]) -> Option<u64> {
    let mut forward = 0_u64;
    let mut reverse = 0_u64;
    for (idx, base) in window.iter().enumerate() {
        let code = base_code(*base)?;
        forward = (forward << 2) | u64::from(code);
        let rc = u64::from(BASE_T - code);
        reverse |= rc << (idx * 2);
    }
    Some(forward.min(reverse))
}

fn base_code(base: u8) -> Option<u8> {
    match base.to_ascii_uppercase() {
        b'A' => Some(BASE_A),
        b'C' => Some(BASE_C),
        b'G' => Some(BASE_G),
        b'T' => Some(BASE_T),
        _ => None,
    }
}

fn mix_u64(mut value: u64) -> u64 {
    value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
    value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
    value ^ (value >> 31)
}

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

    #[test]
    fn rotation_round_trips() -> Result<(), String> {
        let mut vector = (0..128)
            .map(|idx| ((idx as f32 + 1.0) * 0.17).sin())
            .collect::<Vec<_>>();
        l2_normalize(&mut vector);

        let rotated = rotate(&vector, 17)?;
        let decoded = inverse_rotate(&rotated, 17)?;
        let err = mse(&vector, &decoded)?;
        assert!(err < 1.0e-12, "round-trip MSE was {err}");
        Ok(())
    }

    #[test]
    fn more_scalar_bits_reduce_error() -> Result<(), String> {
        let mut vector = (0..256)
            .map(|idx| ((idx as f32 + 3.0) * 0.11).cos())
            .collect::<Vec<_>>();
        l2_normalize(&mut vector);

        let low = QuantizerConfig {
            bits: 2,
            use_qjl_residual: false,
            ..QuantizerConfig::default()
        }
        .encode(&vector)?
        .decode()?;
        let high = QuantizerConfig {
            bits: 5,
            use_qjl_residual: false,
            ..QuantizerConfig::default()
        }
        .encode(&vector)?
        .decode()?;

        assert!(mse(&vector, &high)? < mse(&vector, &low)?);
        Ok(())
    }

    #[test]
    fn packed_codes_round_trip_values() -> Result<(), String> {
        let values = (0..97).map(|idx| (idx % 7) as u8).collect::<Vec<_>>();
        let packed = PackedCodes::encode(&values, 3)?;

        assert!(packed.byte_len() < values.len());
        for (idx, expected) in values.iter().enumerate() {
            assert_eq!(packed.get(idx)?, *expected);
        }
        Ok(())
    }

    #[test]
    fn dna_sketches_are_normalized() -> Result<(), String> {
        let seq = synthetic_dna(512, 9);
        let sketch = dna_kmer_sketch(&seq, 15, 128)?;
        let norm = l2_norm(&sketch);
        assert!((norm - 1.0).abs() < 1.0e-5, "norm was {norm}");
        Ok(())
    }

    #[test]
    fn protein_sketches_are_normalized_and_skip_ambiguous_residues() -> Result<(), String> {
        let sketch = protein_kmer_sketch(b"MKRISTXTTITTTITITTGNGAG", 3, 128)?;
        let norm = l2_norm(&sketch);
        assert!((norm - 1.0).abs() < 1.0e-5, "norm was {norm}");
        Ok(())
    }

    #[test]
    fn parses_fasta_records_and_concatenates_with_separator() -> Result<(), String> {
        let records = parse_fasta_bytes(b">chr1 description\nacgt\nNN\n>chr2\nTTA\n")?;
        assert_eq!(records.len(), 2);
        assert_eq!(records[0].name, "chr1");
        assert_eq!(records[0].bases, b"ACGTNN");
        assert_eq!(records[1].name, "chr2");
        assert_eq!(concatenate_records(&records)?, b"ACGTNNNTTA");
        Ok(())
    }

    #[test]
    fn parses_fastq_records() -> Result<(), String> {
        let records = parse_fastq_bytes(b"@read1 comment\nacgtn\n+\nIIIII\n@read2\nTTA\n+\n###\n")?;
        assert_eq!(records.len(), 2);
        assert_eq!(records[0].name, "read1");
        assert_eq!(records[0].bases, b"ACGTN");
        assert_eq!(records[1].name, "read2");
        assert_eq!(records[1].bases, b"TTA");
        Ok(())
    }

    #[test]
    fn fastq_slice_visitor_can_stop_after_limit() -> Result<(), String> {
        let input = b"@read1\nACGT\n+\nIIII\n@read2\nTTAA\n+\n####\n";
        let mut reader = dino_seq::FastqReader::new(input.as_slice());
        let mut names = Vec::new();
        visit_fastq_slices_with_reader(&mut reader, Some(1), |record| {
            names.push(String::from_utf8(record.name.to_vec()).map_err(|err| err.to_string())?);
            Ok(())
        })?;
        assert_eq!(names, vec!["read1"]);
        Ok(())
    }

    #[test]
    fn parses_protein_fasta_records() -> Result<(), String> {
        let path =
            std::env::temp_dir().join(format!("dino_quant_protein_{}.faa", std::process::id()));
        std::fs::write(&path, b">p1 protein\nmkristX*\n").map_err(|err| err.to_string())?;
        let records = read_protein_fasta_file(&path)?;
        std::fs::remove_file(path).map_err(|err| err.to_string())?;
        assert_eq!(records.len(), 1);
        assert_eq!(records[0].name, "p1");
        assert_eq!(records[0].bases, b"MKRISTX*");
        Ok(())
    }

    #[test]
    fn rejects_invalid_sequence_character() {
        let err = parse_fasta_bytes(b">chr1\nACGTX\n").expect_err("invalid base should fail");
        assert!(err.contains("unsupported sequence character"));
    }

    #[test]
    fn qjl_residual_decodes_finite_vector() -> Result<(), String> {
        let seq = synthetic_dna(2048, 42);
        let vector = dna_kmer_sketch(&seq, 17, 256)?;
        let quantized = QuantizerConfig {
            bits: 3,
            use_qjl_residual: true,
            ..QuantizerConfig::default()
        }
        .encode(&vector)?;
        let decoded = quantized.decode()?;

        assert_eq!(decoded.len(), vector.len());
        assert!(decoded.iter().all(|value| value.is_finite()));
        assert!(quantized.compressed_bits() < vector.len() * 32);
        Ok(())
    }

    #[test]
    fn approximate_dot_matches_decoded_qjl_dot() -> Result<(), String> {
        let reference = dna_kmer_sketch(&synthetic_dna(2048, 42), 17, 256)?;
        let query = dna_kmer_sketch(&mutate_dna(&synthetic_dna(2048, 77), 19), 17, 256)?;
        let quantized = QuantizerConfig {
            bits: 4,
            use_qjl_residual: true,
            ..QuantizerConfig::default()
        }
        .encode(&reference)?;

        let decoded = quantized.decode()?;
        let decoded_dot = dot(&decoded, &query)?;
        let approximate_dot = quantized.approximate_dot_query(&query)?;

        assert!(
            (decoded_dot - approximate_dot).abs() < 1.0e-5,
            "decoded_dot={decoded_dot} approximate_dot={approximate_dot}"
        );
        Ok(())
    }

    #[test]
    fn quantized_vector_snapshot_round_trips_scoring() -> Result<(), String> {
        let reference = dna_kmer_sketch(&synthetic_dna(512, 19), 11, 128)?;
        let query = dna_kmer_sketch(&mutate_dna(&synthetic_dna(512, 19), 7), 11, 128)?;
        let quantized = QuantizerConfig {
            bits: 4,
            use_qjl_residual: true,
            ..QuantizerConfig::default()
        }
        .encode(&reference)?;
        let restored = QuantizedVector::from_snapshot(quantized.snapshot())?;

        let original_score = quantized.approximate_dot_query(&query)?;
        let restored_score = restored.approximate_dot_query(&query)?;
        assert_eq!(original_score.to_bits(), restored_score.to_bits());
        Ok(())
    }

    #[test]
    fn prepared_query_matches_direct_scoring() -> Result<(), String> {
        let reference = dna_kmer_sketch(&synthetic_dna(512, 23), 11, 128)?;
        let query = dna_kmer_sketch(&mutate_dna(&synthetic_dna(512, 23), 5), 11, 128)?;
        let config = QuantizerConfig {
            bits: 4,
            use_qjl_residual: true,
            ..QuantizerConfig::default()
        };
        let quantized = config.encode(&reference)?;
        let prepared = QuantizedVector::prepare_approximate_query(
            &query,
            config.rotation_seed,
            Some(config.qjl_seed),
        )?;

        let direct = quantized.approximate_dot_query(&query)?;
        let reused = quantized.approximate_dot_prepared_query(&prepared)?;
        assert!(
            (direct - reused).abs() < 1.0e-5,
            "direct={direct} reused={reused}"
        );
        Ok(())
    }

    #[test]
    fn compressed_index_finds_mutated_source_window() -> Result<(), String> {
        let reference = synthetic_dna(8192, 0xabc);
        let config = ReferenceIndexConfig {
            k: 15,
            dim: 256,
            window_len: 512,
            stride: 128,
            quantizer: QuantizerConfig {
                bits: 4,
                use_qjl_residual: false,
                ..QuantizerConfig::default()
            },
        };
        let index = ReferenceWindowIndex::build(&reference, config)?;
        let source_start = 2816;
        let source_end = source_start + config.window_len;
        let query = mutate_dna(&reference[source_start..source_end], 41);
        let hits = index.search_sequence(&query, 5)?;

        assert!(
            hits.iter()
                .any(|hit| { intervals_overlap(hit.start, hit.end, source_start, source_end) })
        );
        assert!(index.compression_ratio() > 6.0);
        Ok(())
    }

    #[test]
    fn record_index_reports_target_coordinates() -> Result<(), String> {
        let records = vec![
            SequenceRecord {
                name: "chr1".to_owned(),
                bases: synthetic_dna(1024, 1),
            },
            SequenceRecord {
                name: "chr2".to_owned(),
                bases: synthetic_dna(1024, 2),
            },
        ];
        let config = ReferenceIndexConfig {
            k: 11,
            dim: 128,
            window_len: 256,
            stride: 128,
            quantizer: QuantizerConfig {
                bits: 4,
                use_qjl_residual: false,
                ..QuantizerConfig::default()
            },
        };
        let index = ReferenceWindowIndex::build_records(&records, config)?;
        let query = records[1].bases[256..512].to_vec();
        let hits = index.search_sequence(&query, 3)?;
        assert!(hits.iter().any(|hit| {
            hit.target_name == "chr2"
                && intervals_overlap(hit.target_start, hit.target_end, 256, 512)
        }));
        Ok(())
    }

    #[test]
    fn canonical_kmers_match_reverse_complements() -> Result<(), String> {
        let left = canonical_kmer_code(b"ACGTT").ok_or("left k-mer should be valid")?;
        let right = canonical_kmer_code(b"AACGT").ok_or("right k-mer should be valid")?;
        assert_eq!(left, right);
        Ok(())
    }
}