ringgrid 0.5.3

//! Marker ID decoding from ring sector patterns.
//!
//! Each marker encodes a unique ID via 16 binary sectors (black/white)
//! in the annular band between the inner and outer ring edges. Decoding:
//!
//! 1. Sample pixel intensities along an elliptical arc in the code band.
//! 2. Threshold into a 16-bit binary word.
//! 3. Match against the embedded codebook, trying all 16 cyclic rotations
//!    and selecting the best by Hamming distance with margin.
//!
//! The codebook is generated by `tools/gen_codebook.py` and embedded as
//! compile-time constants in `codebook.rs`.

use super::codebook::{
    CODEBOOK, CODEBOOK_BITS, CODEBOOK_EXTENDED, CODEBOOK_EXTENDED_MIN_CYCLIC_DIST,
    CODEBOOK_EXTENDED_SEED, CODEBOOK_MIN_CYCLIC_DIST, CODEBOOK_SEED,
};

// ── Bit manipulation helpers ───────────────────────────────────────────

/// Count set bits (popcount).
#[inline]
fn popcount(x: u16) -> u8 {
    x.count_ones() as u8
}

/// Cyclic left rotation of a 16-bit word by `k` positions.
#[inline]
pub fn rotate_left_16(word: u16, k: u32) -> u16 {
    word.rotate_left(k % 16)
}

/// Return the canonical form (minimum over all cyclic rotations).
#[cfg_attr(not(test), allow(dead_code))]
pub fn canonical_16(word: u16) -> u16 {
    (0..16).map(|k| rotate_left_16(word, k)).min().unwrap()
}

#[cfg(test)]
#[inline]
fn cyclic_distance(a: u16, b: u16) -> u8 {
    (0u32..CODEBOOK_BITS as u32)
        .map(|k| popcount(a ^ rotate_left_16(b, k)))
        .min()
        .unwrap()
}

// ── Codebook wrapper ───────────────────────────────────────────────────

/// Explicit embedded codebook profile selector.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CodebookProfile {
    /// Shipped baseline profile with stable IDs `0..892`.
    #[default]
    Base,
    /// Additive opt-in profile that appends the remaining valid 16-bit words
    /// without introducing new complement collisions beyond the shipped
    /// baseline.
    Extended,
}

impl CodebookProfile {
    /// Stable profile name for CLI/docs output.
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Base => "base",
            Self::Extended => "extended",
        }
    }
}

/// An embedded codebook for marker ID matching.
pub struct Codebook {
    profile: CodebookProfile,
    words: &'static [u16],
    min_cyclic_dist: usize,
    seed: u64,
}

impl Default for Codebook {
    fn default() -> Self {
        Self::from_profile(CodebookProfile::Base)
    }
}

/// Result of matching an observed word against the codebook.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Match {
    /// Index into the codebook (marker ID).
    pub id: usize,
    /// Cyclic rotation that produced the best match (0..15).
    pub rotation: u8,
    /// Hamming distance to the best-matching codeword (after rotation).
    pub dist: u8,
    /// Margin: second_best_dist − best_dist.
    pub margin: u8,
    /// Confidence heuristic in [0, 1].
    pub confidence: f32,
}

impl Codebook {
    /// Create a codebook from a static slice of codewords.
    #[cfg_attr(not(test), allow(dead_code))]
    pub fn new(words: &'static [u16]) -> Self {
        Self {
            profile: CodebookProfile::Base,
            words,
            min_cyclic_dist: 1,
            seed: 0,
        }
    }

    /// Create one of the embedded codebook profiles.
    pub fn from_profile(profile: CodebookProfile) -> Self {
        match profile {
            CodebookProfile::Base => Self {
                profile,
                words: &CODEBOOK,
                min_cyclic_dist: CODEBOOK_MIN_CYCLIC_DIST,
                seed: CODEBOOK_SEED,
            },
            CodebookProfile::Extended => Self {
                profile,
                words: &CODEBOOK_EXTENDED,
                min_cyclic_dist: CODEBOOK_EXTENDED_MIN_CYCLIC_DIST,
                seed: CODEBOOK_EXTENDED_SEED,
            },
        }
    }

    /// Selected embedded profile.
    pub fn profile(&self) -> CodebookProfile {
        self.profile
    }

    /// Bits per codeword.
    pub fn bits(&self) -> usize {
        CODEBOOK_BITS
    }

    /// Minimum cyclic Hamming distance claimed for this profile.
    pub fn min_cyclic_dist(&self) -> usize {
        self.min_cyclic_dist
    }

    /// Generator seed recorded for this profile.
    pub fn seed(&self) -> u64 {
        self.seed
    }

    /// Number of codewords.
    #[cfg_attr(not(test), allow(dead_code))]
    pub fn len(&self) -> usize {
        self.words.len()
    }

    /// Whether the codebook is empty.
    #[cfg_attr(not(test), allow(dead_code))]
    pub fn is_empty(&self) -> bool {
        self.words.is_empty()
    }

    /// Get the codeword for a given marker ID.
    #[cfg_attr(not(test), allow(dead_code))]
    pub fn word(&self, id: usize) -> Option<u16> {
        self.words.get(id).copied()
    }

    /// Match an observed 16-bit word against the codebook.
    ///
    /// Tries all 16 cyclic rotations of each codeword and returns the
    /// best match. Confidence combines distance and margin:
    ///   confidence = clamp(1 − dist/6) * clamp(margin / profile.min_cyclic_dist)
    ///
    /// The margin denominator equals the codebook minimum cyclic Hamming distance,
    /// so a perfect decode (dist=0) with the minimum attainable margin scores 1.0.
    pub fn match_word(&self, obs: u16) -> Match {
        let mut best_id = 0usize;
        let mut best_rot = 0u8;
        let mut best_dist = u8::MAX;
        let mut second_dist = u8::MAX;

        for (id, &cw) in self.words.iter().enumerate() {
            for k in 0u32..self.bits() as u32 {
                let rotated = rotate_left_16(cw, k);
                let d = popcount(obs ^ rotated);
                if d < best_dist {
                    second_dist = best_dist;
                    best_dist = d;
                    best_id = id;
                    best_rot = k as u8;
                } else if d < second_dist {
                    second_dist = d;
                }
            }
        }

        let margin = second_dist.saturating_sub(best_dist);
        let conf_dist = (1.0 - best_dist as f32 / 6.0).clamp(0.0, 1.0);
        let conf_margin = (margin as f32 / self.min_cyclic_dist as f32).clamp(0.0, 1.0);
        let confidence = conf_dist * conf_margin;

        Match {
            id: best_id,
            rotation: best_rot,
            dist: best_dist,
            margin,
            confidence,
        }
    }
}

// ── Tests ──────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::super::codebook::{
        CODEBOOK_EXTENDED, CODEBOOK_EXTENDED_EXTENSION_N, CODEBOOK_EXTENDED_MIN_CYCLIC_DIST,
        CODEBOOK_EXTENDED_N, CODEBOOK_MIN_CYCLIC_DIST, CODEBOOK_N,
    };
    use super::*;
    use rand::RngExt;
    use rand::prelude::*;
    use std::collections::HashSet;

    fn assert_no_codeword_is_rotationally_symmetric(words: &[u16]) {
        for (i, &w) in words.iter().enumerate() {
            for k in 1u32..16 {
                assert_ne!(
                    rotate_left_16(w, k),
                    w,
                    "codeword {} (0x{:04X}) is rotationally symmetric at k={}",
                    i,
                    w,
                    k
                );
            }
        }
    }

    fn assert_pairwise_uniqueness_under_rotation(words: &[u16]) {
        let mut canonical_words = HashSet::with_capacity(words.len());
        for (i, &w) in words.iter().enumerate() {
            assert!(
                canonical_words.insert(canonical_16(w)),
                "codeword {} (0x{:04X}) collides under rotation with another entry",
                i,
                w
            );
        }
    }

    #[test]
    fn test_no_codeword_is_rotationally_symmetric() {
        assert_no_codeword_is_rotationally_symmetric(&CODEBOOK);
        assert_no_codeword_is_rotationally_symmetric(&CODEBOOK_EXTENDED);
    }

    #[test]
    fn test_pairwise_uniqueness_under_rotation() {
        assert_pairwise_uniqueness_under_rotation(&CODEBOOK);
        assert_pairwise_uniqueness_under_rotation(&CODEBOOK_EXTENDED);
    }

    #[test]
    fn test_pairwise_min_cyclic_distance() {
        // Verify a sample of pairs meets the claimed minimum distance
        let n = CODEBOOK.len();
        let mut rng = StdRng::seed_from_u64(123);
        let pairs = 20_000.min(n * (n - 1) / 2);
        let mut observed_min = 16u8;

        for _ in 0..pairs {
            let i = rng.random_range(0..n);
            let j = rng.random_range(0..n);
            if i == j {
                continue;
            }
            let d = cyclic_distance(CODEBOOK[i], CODEBOOK[j]);
            if d < observed_min {
                observed_min = d;
            }
        }

        assert!(
            observed_min >= CODEBOOK_MIN_CYCLIC_DIST as u8,
            "observed min cyclic dist {} < claimed {}",
            observed_min,
            CODEBOOK_MIN_CYCLIC_DIST
        );
    }

    #[test]
    fn test_default_profile_matches_shipped_baseline() {
        let cb = Codebook::default();
        assert_eq!(cb.profile(), CodebookProfile::Base);
        assert_eq!(cb.len(), CODEBOOK_N);
        assert_eq!(cb.min_cyclic_dist(), CODEBOOK_MIN_CYCLIC_DIST);
        assert_eq!(cb.seed(), CODEBOOK_SEED);
    }

    #[test]
    fn test_extended_profile_appends_base_prefix() {
        let cb = Codebook::from_profile(CodebookProfile::Extended);
        assert_eq!(cb.profile(), CodebookProfile::Extended);
        assert_eq!(cb.len(), CODEBOOK_EXTENDED_N);
        assert_eq!(cb.min_cyclic_dist(), CODEBOOK_EXTENDED_MIN_CYCLIC_DIST);
        assert_eq!(&CODEBOOK_EXTENDED[..CODEBOOK_N], &CODEBOOK);
        assert_eq!(
            CODEBOOK_EXTENDED_EXTENSION_N,
            CODEBOOK_EXTENDED_N - CODEBOOK_N
        );
    }

    #[test]
    fn test_extended_profile_has_distance_one_witness() {
        let appended = &CODEBOOK_EXTENDED[CODEBOOK_N..];
        assert!(
            appended
                .iter()
                .any(|&ext| CODEBOOK.iter().any(|&base| cyclic_distance(ext, base) == 1)),
            "extended profile should witness the weaker min distance of 1"
        );
    }

    #[test]
    fn test_extended_profile_appendix_excludes_new_complement_collisions() {
        let canonical_words: HashSet<u16> =
            CODEBOOK_EXTENDED.iter().map(|&w| canonical_16(w)).collect();
        for (offset, &w) in CODEBOOK_EXTENDED[CODEBOOK_N..].iter().enumerate() {
            let canonical = canonical_16(w);
            let complement = canonical_16(!w);
            if complement == canonical {
                continue;
            }
            assert!(
                !canonical_words.contains(&complement),
                "appended codeword {} (0x{:04X}) collides with complement class 0x{:04X}",
                CODEBOOK_N + offset,
                w,
                complement
            );
        }
    }

    #[test]
    fn test_extended_profile_exact_match_for_first_appended_id() {
        let cb = Codebook::from_profile(CodebookProfile::Extended);
        let w = cb
            .word(CODEBOOK_N)
            .expect("extended profile should expose appended IDs");
        let m = cb.match_word(w);
        assert_eq!(m.id, CODEBOOK_N);
        assert_eq!(m.dist, 0);
        assert_eq!(m.rotation, 0);
    }

    #[test]
    fn test_exact_match() {
        let cb = Codebook::default();
        // Every codeword should match itself exactly
        for id in 0..cb.len().min(100) {
            let w = cb.word(id).unwrap();
            let m = cb.match_word(w);
            assert_eq!(m.id, id, "exact match failed for id {}", id);
            assert_eq!(m.dist, 0);
            assert_eq!(m.rotation, 0);
        }
    }

    #[test]
    fn test_match_with_rotation() {
        let cb = Codebook::default();
        let mut rng = StdRng::seed_from_u64(77);
        for _ in 0..100 {
            let id = rng.random_range(0..cb.len());
            let w = cb.word(id).unwrap();
            let rot = rng.random_range(0u32..16);
            let rotated = rotate_left_16(w, rot);
            let m = cb.match_word(rotated);
            assert_eq!(m.id, id, "rotation match failed for id {} rot {}", id, rot);
            assert_eq!(m.dist, 0);
        }
    }

    #[test]
    fn test_match_with_1_bit_flip() {
        // With min_cyclic_dist=2, a 1-bit flip guarantees the corrupted word
        // is closer to the true codeword (dist=1) than to any other (dist>=2).
        // However, another codeword *rotated* might also be at dist=1.
        // So we check that the best match distance is 1 and that most decode
        // correctly.
        let cb = Codebook::default();
        let mut rng = StdRng::seed_from_u64(99);
        let mut correct = 0;
        let total = 200;

        for _ in 0..total {
            let id = rng.random_range(0..cb.len());
            let w = cb.word(id).unwrap();
            let rot = rng.random_range(0u32..16);
            let mut obs = rotate_left_16(w, rot);

            let bit = rng.random_range(0..16);
            obs ^= 1u16 << bit;

            let m = cb.match_word(obs);
            assert!(
                m.dist <= 1,
                "1-bit flip should give dist <= 1, got {}",
                m.dist
            );
            if m.id == id {
                correct += 1;
            }
        }

        let rate = correct as f64 / total as f64;
        eprintln!("1-bit flip correct rate: {:.1}%", rate * 100.0);
        // With min_cyclic_dist=2, 1-bit error correction is not guaranteed.
        // We just verify the machinery doesn't panic and returns dist <= 1.
    }

    #[test]
    fn test_match_with_2_bit_flips() {
        // With min_cyclic_dist=2, 2-bit flips can land on another codeword
        // (dist=2 = min_dist). We just verify no panics and report the rate.
        let cb = Codebook::default();
        let mut rng = StdRng::seed_from_u64(2024);
        let mut correct = 0;
        let total = 200;

        for _ in 0..total {
            let id = rng.random_range(0..cb.len());
            let w = cb.word(id).unwrap();
            let rot = rng.random_range(0u32..16);
            let mut obs = rotate_left_16(w, rot);

            let bit1 = rng.random_range(0u16..16);
            let mut bit2 = rng.random_range(0u16..15);
            if bit2 >= bit1 {
                bit2 += 1;
            }
            obs ^= (1u16 << bit1) | (1u16 << bit2);

            let m = cb.match_word(obs);
            assert!(
                m.dist <= 2,
                "2-bit flip should give dist <= 2, got {}",
                m.dist
            );
            if m.id == id {
                correct += 1;
            }
        }

        // Just log the rate; with dist=2, 2-bit flips are at the correction limit
        let rate = correct as f64 / total as f64;
        eprintln!("2-bit flip correct rate: {:.1}%", rate * 100.0);
        // No hard assertion on rate — this is informational
    }

    #[test]
    fn test_confidence_heuristic() {
        let cb = Codebook::default();
        let w = cb.word(0).unwrap();

        // Exact match: high confidence
        let m = cb.match_word(w);
        assert!(
            m.confidence > 0.0,
            "exact match should have positive confidence"
        );

        // Far-away word (all bits flipped): low distance to complement
        let flipped = !w;
        let m2 = cb.match_word(flipped);
        // The complement might match some other codeword well, but generally
        // with large distance the confidence should be lower
        let _ = m2; // just ensure no panic
    }

    #[test]
    fn test_canonical() {
        assert_eq!(canonical_16(0x0001), 0x0001);
        assert_eq!(canonical_16(0x8000), 0x0001); // rotation of 0x0001
        assert_eq!(canonical_16(0x0003), 0x0003);
        assert_eq!(canonical_16(0xC000), 0x0003); // rotation of 0x0003
    }
}