scirs2-vision 0.5.1

//! BRIEF (Binary Robust Independent Elementary Features) descriptor.
//!
//! Computes compact binary descriptors by comparing pairs of smoothed pixel
//! intensities within a patch around a keypoint.  The result is a packed
//! bitset stored as a `Vec<u64>` for efficient Hamming-distance computation.
//!
//! Reference: Calonder et al., "BRIEF: Binary Robust Independent Elementary
//! Features", ECCV 2010.

use crate::error::VisionError;

// ─── Configuration ────────────────────────────────────────────────────────────

/// Sampling pattern for BRIEF test pairs.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BRIEFPattern {
    /// Pairs drawn uniformly from [-patch_size/2, patch_size/2].
    Uniform,
    /// Pairs drawn from N(0, (patch_size/5)²) (isotropic Gaussian).
    Gaussian,
}

/// Configuration for BRIEF descriptor computation.
#[derive(Debug, Clone)]
pub struct BRIEFConfig {
    /// Number of binary tests (descriptor length in bits). Should be a multiple of 64.
    pub n_bits: usize,
    /// Side length of the patch (pixels). Must be odd.
    pub patch_size: usize,
    /// Distribution used to generate test pairs.
    pub pattern: BRIEFPattern,
}

impl Default for BRIEFConfig {
    fn default() -> Self {
        Self {
            n_bits: 256,
            patch_size: 31,
            pattern: BRIEFPattern::Gaussian,
        }
    }
}

// ─── Descriptor type ─────────────────────────────────────────────────────────

/// BRIEF binary descriptor stored as packed 64-bit words.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BRIEFDescriptor {
    /// Packed bits.  `bits.len() == n_bits / 64` (rounded up).
    pub bits: Vec<u64>,
    /// Number of valid bits (may differ from `bits.len() * 64`).
    pub n_bits: usize,
}

// ─── Test-pair generation ────────────────────────────────────────────────────

/// Minimal xorshift64 PRNG seeded with `seed`.
fn xorshift64(state: &mut u64) -> u64 {
    *state ^= *state << 13;
    *state ^= *state >> 7;
    *state ^= *state << 17;
    *state
}

/// Sample a value in `[-half, half]` uniformly using the provided PRNG state.
fn sample_uniform(state: &mut u64, half: i32) -> i32 {
    let range = (2 * half + 1) as u64;
    let v = xorshift64(state) % range;
    v as i32 - half
}

/// Sample from N(0, sigma²) using Box-Muller transform.
fn sample_gaussian(state: &mut u64, sigma: f64) -> i32 {
    // Box-Muller needs two uniform [0,1) samples.
    const MAX: f64 = u64::MAX as f64;
    let u1 = (xorshift64(state) as f64 + 1.0) / (MAX + 2.0);
    let u2 = (xorshift64(state) as f64 + 1.0) / (MAX + 2.0);
    let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
    (z * sigma).round() as i32
}

/// Generate `n_bits` pixel-pair test offsets for BRIEF.
///
/// Each entry is `((dy_p, dx_p), (dy_q, dx_q))` relative to the keypoint centre.
/// Offsets are clamped to `[-patch_size/2, patch_size/2]`.
pub fn generate_test_pairs(
    n_bits: usize,
    patch_size: usize,
    pattern: BRIEFPattern,
    seed: u64,
) -> Vec<((i32, i32), (i32, i32))> {
    let half = (patch_size / 2) as i32;
    let sigma = patch_size as f64 / 5.0;
    let mut state = if seed == 0 { 0xdeadbeef_u64 } else { seed };
    let mut pairs = Vec::with_capacity(n_bits);

    for _ in 0..n_bits {
        let (py, px, qy, qx) = match pattern {
            BRIEFPattern::Uniform => (
                sample_uniform(&mut state, half),
                sample_uniform(&mut state, half),
                sample_uniform(&mut state, half),
                sample_uniform(&mut state, half),
            ),
            BRIEFPattern::Gaussian => (
                sample_gaussian(&mut state, sigma).clamp(-half, half),
                sample_gaussian(&mut state, sigma).clamp(-half, half),
                sample_gaussian(&mut state, sigma).clamp(-half, half),
                sample_gaussian(&mut state, sigma).clamp(-half, half),
            ),
        };
        pairs.push(((py, px), (qy, qx)));
    }
    pairs
}

// ─── Descriptor computation ──────────────────────────────────────────────────

/// Look up a (possibly clamped) pixel in `image`.
#[inline]
fn get_pixel(image: &[Vec<f32>], r: i32, c: i32, rows: usize, cols: usize) -> f32 {
    let rr = r.clamp(0, rows as i32 - 1) as usize;
    let cc = c.clamp(0, cols as i32 - 1) as usize;
    image[rr][cc]
}

/// Compute a BRIEF descriptor for the patch centred at `keypoint`.
///
/// `pairs` must have been generated by [`generate_test_pairs`] with the same
/// `n_bits` as `config.n_bits`.
///
/// # Errors
/// Returns [`VisionError`] if the image is empty.
#[allow(clippy::type_complexity)]
pub fn compute_brief(
    image: &[Vec<f32>],
    keypoint: (usize, usize),
    pairs: &[((i32, i32), (i32, i32))],
) -> Result<BRIEFDescriptor, VisionError> {
    let rows = image.len();
    let cols = image.first().map_or(0, |r| r.len());
    if rows == 0 || cols == 0 {
        return Err(VisionError::InvalidInput("Empty image".into()));
    }

    let ky = keypoint.0 as i32;
    let kx = keypoint.1 as i32;
    let n_bits = pairs.len();
    let n_words = n_bits.div_ceil(64);
    let mut bits = vec![0_u64; n_words];

    for (i, &((dy_p, dx_p), (dy_q, dx_q))) in pairs.iter().enumerate() {
        let p = get_pixel(image, ky + dy_p, kx + dx_p, rows, cols);
        let q = get_pixel(image, ky + dy_q, kx + dx_q, rows, cols);
        if p < q {
            bits[i / 64] |= 1_u64 << (i % 64);
        }
    }

    Ok(BRIEFDescriptor { bits, n_bits })
}

// ─── Distance ────────────────────────────────────────────────────────────────

/// Compute the Hamming distance between two BRIEF descriptors.
///
/// Descriptors must have been generated with the same `n_bits`; if they
/// differ in word count the shorter one is zero-padded implicitly.
pub fn hamming_distance(a: &BRIEFDescriptor, b: &BRIEFDescriptor) -> u32 {
    let len = a.bits.len().max(b.bits.len());
    (0..len)
        .map(|i| {
            let wa = a.bits.get(i).copied().unwrap_or(0);
            let wb = b.bits.get(i).copied().unwrap_or(0);
            (wa ^ wb).count_ones()
        })
        .sum()
}

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

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

    fn flat_image(rows: usize, cols: usize, val: f32) -> Vec<Vec<f32>> {
        vec![vec![val; cols]; rows]
    }

    fn gradient_image(rows: usize, cols: usize) -> Vec<Vec<f32>> {
        (0..rows)
            .map(|r| {
                (0..cols)
                    .map(|c| (r * cols + c) as f32 / (rows * cols) as f32)
                    .collect()
            })
            .collect()
    }

    #[test]
    fn test_generate_pairs_count() {
        let pairs = generate_test_pairs(256, 31, BRIEFPattern::Gaussian, 42);
        assert_eq!(pairs.len(), 256);
    }

    #[test]
    fn test_generate_pairs_bounds_uniform() {
        let patch_size = 31_usize;
        let half = (patch_size / 2) as i32;
        let pairs = generate_test_pairs(256, patch_size, BRIEFPattern::Uniform, 7);
        for ((dy_p, dx_p), (dy_q, dx_q)) in &pairs {
            assert!((-half..=half).contains(dy_p));
            assert!((-half..=half).contains(dx_p));
            assert!((-half..=half).contains(dy_q));
            assert!((-half..=half).contains(dx_q));
        }
    }

    #[test]
    fn test_flat_image_zero_descriptor() {
        // All pixels equal → all tests give 0 → all bits 0.
        let img = flat_image(64, 64, 0.5);
        let pairs = generate_test_pairs(256, 31, BRIEFPattern::Gaussian, 1);
        let desc = compute_brief(&img, (32, 32), &pairs)
            .expect("compute_brief should succeed on valid image");
        assert!(desc.bits.iter().all(|&w| w == 0));
    }

    #[test]
    fn test_descriptor_word_count() {
        let img = gradient_image(64, 64);
        let pairs = generate_test_pairs(256, 31, BRIEFPattern::Gaussian, 1);
        let desc = compute_brief(&img, (32, 32), &pairs)
            .expect("compute_brief should succeed on valid image");
        assert_eq!(desc.bits.len(), 4); // 256 / 64
        assert_eq!(desc.n_bits, 256);
    }

    #[test]
    fn test_hamming_distance_self() {
        let img = gradient_image(64, 64);
        let pairs = generate_test_pairs(128, 31, BRIEFPattern::Gaussian, 99);
        let d = compute_brief(&img, (32, 32), &pairs)
            .expect("compute_brief should succeed on valid image");
        assert_eq!(hamming_distance(&d, &d), 0);
    }

    #[test]
    fn test_hamming_distance_symmetry() {
        let img = gradient_image(64, 64);
        let pairs = generate_test_pairs(128, 31, BRIEFPattern::Uniform, 55);
        let d1 = compute_brief(&img, (20, 20), &pairs)
            .expect("compute_brief should succeed at keypoint (20,20)");
        let d2 = compute_brief(&img, (40, 40), &pairs)
            .expect("compute_brief should succeed at keypoint (40,40)");
        assert_eq!(hamming_distance(&d1, &d2), hamming_distance(&d2, &d1));
    }

    #[test]
    fn test_hamming_complement() {
        // Two descriptors that differ in every bit should have distance n_bits.
        let a = BRIEFDescriptor {
            bits: vec![0x0000_0000_0000_0000_u64; 4],
            n_bits: 256,
        };
        let b = BRIEFDescriptor {
            bits: vec![0xFFFF_FFFF_FFFF_FFFF_u64; 4],
            n_bits: 256,
        };
        assert_eq!(hamming_distance(&a, &b), 256);
    }
}