//! Salient-point detector and descriptor, built for hashing experiments.
//!
//! The detector is SIFT-like: extrema in a Difference-of-Gaussians scale
//! space, filtered for contrast and edge responses. The descriptor is
//! SURF-like: a 4x4 grid of subregions around the keypoint, each summarized
//! by four gradient sums, 64 floats total. Every descriptor is measured in
//! the keypoint's own scale and orientation frame, which is where the scale
//! and rotation invariance comes from.
//!
//! Descriptors are canonicalized under horizontal mirroring: a flipped patch
//! produces the same 64 floats, so hashes built from them inherit flip
//! invariance the way the DCT hash does. See `mirror_descriptor` for the
//! exact symmetry.
//!
//! Not wired into the Descriptor trait or the stores. The salientlab binary
//! runs the experiments.

use image::DynamicImage;
use std::f32::consts::PI;

const ORI_BINS: usize = 36;

/// Descriptor sampling grid: DESC_SAMPLES x DESC_SAMPLES points spanning
/// +-DESC_HALF_EXTENT keypoint sigmas, grouped into DESC_GRID x DESC_GRID
/// subregions. These define the exact window the descriptor reads, so the
/// browser demo draws the same box by calling Keypoint::window_corners.
const DESC_SAMPLES: usize = 20;
const DESC_HALF_EXTENT: f32 = 9.5;
pub const DESC_GRID: usize = 4;

/// Grayscale f32 image, values in [0, 1]
pub struct GrayF32 {
    w: usize,
    h: usize,
    data: Vec<f32>,
}

impl GrayF32 {
    pub fn from_image(img: &DynamicImage) -> GrayF32 {
        let luma = img.to_luma32f();
        GrayF32 {
            w: luma.width() as usize,
            h: luma.height() as usize,
            data: luma.into_raw(),
        }
    }

    fn get(&self, x: isize, y: isize) -> f32 {
        let x = x.clamp(0, self.w as isize - 1) as usize;
        let y = y.clamp(0, self.h as isize - 1) as usize;
        self.data[y * self.w + x]
    }

    fn bilinear(&self, x: f32, y: f32) -> f32 {
        let x0 = x.floor();
        let y0 = y.floor();
        let fx = x - x0;
        let fy = y - y0;
        let (xi, yi) = (x0 as isize, y0 as isize);
        let v00 = self.get(xi, yi);
        let v10 = self.get(xi + 1, yi);
        let v01 = self.get(xi, yi + 1);
        let v11 = self.get(xi + 1, yi + 1);
        v00 * (1.0 - fx) * (1.0 - fy)
            + v10 * fx * (1.0 - fy)
            + v01 * (1.0 - fx) * fy
            + v11 * fx * fy
    }

    /// Image gradient at a fractional position, central differences
    fn gradient(&self, x: f32, y: f32) -> (f32, f32) {
        (
            (self.bilinear(x + 1.0, y) - self.bilinear(x - 1.0, y)) * 0.5,
            (self.bilinear(x, y + 1.0) - self.bilinear(x, y - 1.0)) * 0.5,
        )
    }

    /// Separable gaussian blur, borders clamped
    fn gauss_blur(&self, sigma: f32) -> GrayF32 {
        let radius = (sigma * 3.0).ceil().max(1.0) as isize;
        let mut kernel = Vec::with_capacity(2 * radius as usize + 1);
        for i in -radius..=radius {
            kernel.push((-(i * i) as f32 / (2.0 * sigma * sigma)).exp());
        }
        let sum: f32 = kernel.iter().sum();
        for k in kernel.iter_mut() {
            *k /= sum;
        }

        let mut tmp = vec![0.0f32; self.w * self.h];
        for y in 0..self.h {
            for x in 0..self.w {
                let mut acc = 0.0;
                for (ki, k) in kernel.iter().enumerate() {
                    acc += k * self.get(x as isize + ki as isize - radius, y as isize);
                }
                tmp[y * self.w + x] = acc;
            }
        }
        let tmp = GrayF32 { w: self.w, h: self.h, data: tmp };
        let mut out = vec![0.0f32; self.w * self.h];
        for y in 0..self.h {
            for x in 0..self.w {
                let mut acc = 0.0;
                for (ki, k) in kernel.iter().enumerate() {
                    acc += k * tmp.get(x as isize, y as isize + ki as isize - radius);
                }
                out[y * self.w + x] = acc;
            }
        }
        GrayF32 { w: self.w, h: self.h, data: out }
    }

    /// Every other pixel, for the next octave
    fn half(&self) -> GrayF32 {
        let w = (self.w / 2).max(1);
        let h = (self.h / 2).max(1);
        let mut data = Vec::with_capacity(w * h);
        for y in 0..h {
            for x in 0..w {
                data.push(self.data[(y * 2) * self.w + x * 2]);
            }
        }
        GrayF32 { w, h, data }
    }

    fn sub(a: &GrayF32, b: &GrayF32) -> GrayF32 {
        let data = a.data.iter().zip(&b.data).map(|(a, b)| a - b).collect();
        GrayF32 { w: a.w, h: a.h, data }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Keypoint {
    /// Position in original image coordinates
    pub x: f32,
    pub y: f32,
    /// Scale in original image coordinates
    pub sigma: f32,
    /// Dominant gradient orientation, radians
    pub angle: f32,
    /// Absolute DoG value at the extremum
    pub response: f32,
}

impl Keypoint {
    /// Map a point from the keypoint's oriented frame (u along the
    /// orientation, v perpendicular) into image coordinates. Same rotation
    /// the descriptor uses to place its samples.
    fn to_image(&self, u: f32, v: f32) -> (f32, f32) {
        let (sin, cos) = self.angle.sin_cos();
        (self.x + u * cos - v * sin, self.y + u * sin + v * cos)
    }

    /// The four corners of the oriented descriptor window in image
    /// coordinates, from the (-u,-v) corner counter-clockwise. For drawing
    /// the region the descriptor summarizes.
    pub fn window_corners(&self) -> [(f32, f32); 4] {
        let e = DESC_HALF_EXTENT * self.sigma;
        [(-e, -e), (e, -e), (e, e), (-e, e)].map(|(u, v)| self.to_image(u, v))
    }

    /// Centres of the DESC_GRID x DESC_GRID subregions in image coordinates,
    /// in descriptor order (row = perpendicular to orientation, column =
    /// along it), so a caller can line them up with the 64-float descriptor.
    pub fn subregion_centers(&self) -> [(f32, f32); DESC_GRID * DESC_GRID] {
        let cell = DESC_SAMPLES / DESC_GRID;
        let mut out = [(0.0, 0.0); DESC_GRID * DESC_GRID];
        for j in 0..DESC_GRID {
            for i in 0..DESC_GRID {
                let su = (i * cell + cell / 2) as f32 - DESC_HALF_EXTENT;
                let sv = (j * cell + cell / 2) as f32 - DESC_HALF_EXTENT;
                out[j * DESC_GRID + i] = self.to_image(su * self.sigma, sv * self.sigma);
            }
        }
        out
    }
}

#[derive(Clone)]
pub struct Feature {
    pub kp: Keypoint,
    pub desc: [f32; 64],
}

pub struct Params {
    /// DoG scale samples per octave
    pub intervals: usize,
    /// Blur of the first level of each octave
    pub base_sigma: f32,
    /// Minimum absolute DoG value for a keypoint (image values are 0..1)
    pub contrast_threshold: f32,
    /// Maximum principal curvature ratio, rejects points on straight edges
    pub edge_ratio: f32,
    /// Keep only the strongest keypoints
    pub max_features: usize,
    /// Canonicalize descriptors under horizontal mirroring
    pub flip_canonical: bool,
}

impl Default for Params {
    fn default() -> Self {
        Params {
            intervals: 3,
            base_sigma: 1.6,
            contrast_threshold: 0.01,
            edge_ratio: 10.0,
            max_features: 100,
            flip_canonical: true,
        }
    }
}

struct RawKp {
    octave: usize,
    level: usize,
    x: usize,
    y: usize,
    bx: f32,
    by: f32,
    bsigma: f32,
    response: f32,
}

pub struct Salient {
    pub params: Params,
}

impl Default for Salient {
    fn default() -> Self {
        Self::new()
    }
}

impl Salient {
    pub fn new() -> Salient {
        Salient { params: Params::default() }
    }

    fn k(&self) -> f32 {
        2f32.powf(1.0 / self.params.intervals as f32)
    }

    /// Gaussian pyramid: per octave `intervals + 3` blur levels, so the DoG
    /// stack has `intervals` usable middle layers
    fn pyramid(&self, img: &DynamicImage) -> Vec<Vec<GrayF32>> {
        let s = self.params.intervals;
        let k = self.k();
        let sigma0 = self.params.base_sigma;
        // assume the input carries sigma 0.5, lift it to base_sigma
        let gray = GrayF32::from_image(img);
        let mut base = gray.gauss_blur((sigma0 * sigma0 - 0.25).max(0.01).sqrt());

        let mut octaves = Vec::new();
        while base.w.min(base.h) >= 24 && octaves.len() < 6 {
            let mut levels = vec![base];
            for i in 1..s + 3 {
                let sig_prev = sigma0 * k.powi(i as i32 - 1);
                let sig_total = sigma0 * k.powi(i as i32);
                let sig_diff = (sig_total * sig_total - sig_prev * sig_prev).sqrt();
                levels.push(levels[i - 1].gauss_blur(sig_diff));
            }
            // levels[s] carries 2 * base_sigma: halved it seeds the next octave
            base = levels[s].half();
            octaves.push(levels);
        }
        octaves
    }

    fn detect(&self, pyramid: &[Vec<GrayF32>]) -> Vec<RawKp> {
        let k = self.k();
        let mut raw = Vec::new();
        for (o, levels) in pyramid.iter().enumerate() {
            let dogs: Vec<GrayF32> = levels
                .windows(2)
                .map(|w| GrayF32::sub(&w[1], &w[0]))
                .collect();
            let scale_up = (1usize << o) as f32;
            for l in 1..dogs.len() - 1 {
                let d = &dogs[l];
                for y in 1..d.h - 1 {
                    for x in 1..d.w - 1 {
                        let v = d.data[y * d.w + x];
                        if v.abs() < self.params.contrast_threshold {
                            continue;
                        }
                        if !is_extremum(&dogs, l, x, y, v) {
                            continue;
                        }
                        if is_edge(d, x, y, self.params.edge_ratio) {
                            continue;
                        }
                        let osigma = self.params.base_sigma * k.powi(l as i32);
                        raw.push(RawKp {
                            octave: o,
                            level: l,
                            x,
                            y,
                            bx: x as f32 * scale_up,
                            by: y as f32 * scale_up,
                            bsigma: osigma * scale_up,
                            response: v.abs(),
                        });
                    }
                }
            }
        }
        raw
    }

    /// Detect, orient and describe the strongest keypoints of an image
    pub fn features(&self, img: &DynamicImage) -> Vec<Feature> {
        let pyramid = self.pyramid(img);
        let mut raw = self.detect(&pyramid);
        raw.sort_by(|a, b| b.response.total_cmp(&a.response));

        // greedy dedupe: extrema often fire in adjacent scale layers
        let mut picked: Vec<RawKp> = Vec::new();
        for kp in raw {
            if picked.len() >= self.params.max_features {
                break;
            }
            let dup = picked.iter().any(|p| {
                let (dx, dy) = (p.bx - kp.bx, p.by - kp.by);
                let ratio = (p.bsigma / kp.bsigma).max(kp.bsigma / p.bsigma);
                dx * dx + dy * dy < 4.0 && ratio < 1.6
            });
            if !dup {
                picked.push(kp);
            }
        }

        picked
            .iter()
            .map(|kp| {
                let img = &pyramid[kp.octave][kp.level];
                let osigma = self.params.base_sigma * self.k().powi(kp.level as i32);
                let angle = orientation(img, kp.x, kp.y, osigma);
                let mut desc = describe(img, kp.x as f32, kp.y as f32, osigma, angle);
                if self.params.flip_canonical {
                    desc = canonicalize(desc);
                }
                Feature {
                    kp: Keypoint {
                        x: kp.bx,
                        y: kp.by,
                        sigma: kp.bsigma,
                        angle,
                        response: kp.response,
                    },
                    desc,
                }
            })
            .collect()
    }
}

/// Strict local extremum over the 26 neighbors in space and scale
fn is_extremum(dogs: &[GrayF32], l: usize, x: usize, y: usize, v: f32) -> bool {
    for layer in &dogs[l - 1..=l + 1] {
        for dy in -1isize..=1 {
            for dx in -1isize..=1 {
                let n = layer.get(x as isize + dx, y as isize + dy);
                if std::ptr::eq(layer, &dogs[l]) && dx == 0 && dy == 0 {
                    continue;
                }
                if (v > 0.0 && n >= v) || (v < 0.0 && n <= v) {
                    return false;
                }
            }
        }
    }
    true
}

/// Principal curvature ratio test on the 2x2 spatial Hessian of the DoG,
/// rejects points that sit on an edge and slide along it
fn is_edge(d: &GrayF32, x: usize, y: usize, r: f32) -> bool {
    let (x, y) = (x as isize, y as isize);
    let v = d.get(x, y);
    let dxx = d.get(x + 1, y) - 2.0 * v + d.get(x - 1, y);
    let dyy = d.get(x, y + 1) - 2.0 * v + d.get(x, y - 1);
    let dxy = (d.get(x + 1, y + 1) - d.get(x + 1, y - 1) - d.get(x - 1, y + 1)
        + d.get(x - 1, y - 1))
        / 4.0;
    let tr = dxx + dyy;
    let det = dxx * dyy - dxy * dxy;
    det <= 0.0 || tr * tr * r >= (r + 1.0) * (r + 1.0) * det
}

/// Dominant gradient orientation: 36-bin histogram of gradient angles in a
/// gaussian-weighted neighborhood, smoothed, peak refined by a parabola
fn orientation(img: &GrayF32, x: usize, y: usize, sigma: f32) -> f32 {
    let sig_w = 1.5 * sigma;
    let radius = (3.0 * sig_w).round() as isize;
    let mut hist = [0f32; ORI_BINS];
    for dy in -radius..=radius {
        for dx in -radius..=radius {
            let (px, py) = (x as isize + dx, y as isize + dy);
            let gx = (img.get(px + 1, py) - img.get(px - 1, py)) * 0.5;
            let gy = (img.get(px, py + 1) - img.get(px, py - 1)) * 0.5;
            let mag = (gx * gx + gy * gy).sqrt();
            if mag == 0.0 {
                continue;
            }
            let mut phi = gy.atan2(gx);
            if phi < 0.0 {
                phi += 2.0 * PI;
            }
            let w = (-((dx * dx + dy * dy) as f32) / (2.0 * sig_w * sig_w)).exp();
            let bin = ((phi / (2.0 * PI) * ORI_BINS as f32) as usize).min(ORI_BINS - 1);
            hist[bin] += w * mag;
        }
    }
    // two passes of circular [1 2 1] smoothing
    for _ in 0..2 {
        let orig = hist;
        for b in 0..ORI_BINS {
            let l = orig[(b + ORI_BINS - 1) % ORI_BINS];
            let r = orig[(b + 1) % ORI_BINS];
            hist[b] = (l + 2.0 * orig[b] + r) / 4.0;
        }
    }
    let peak = (0..ORI_BINS)
        .max_by(|a, b| hist[*a].total_cmp(&hist[*b]))
        .unwrap();
    let l = hist[(peak + ORI_BINS - 1) % ORI_BINS];
    let c = hist[peak];
    let r = hist[(peak + 1) % ORI_BINS];
    let denom = l - 2.0 * c + r;
    let delta = if denom.abs() > 1e-12 { 0.5 * (l - r) / denom } else { 0.0 };
    (peak as f32 + 0.5 + delta) * 2.0 * PI / ORI_BINS as f32
}

/// SURF-style descriptor: 20x20 samples spaced `sigma` apart, rotated to the
/// keypoint orientation, grouped into 4x4 subregions. Each subregion gets
/// [sum du, sum dv, sum |du|, sum |dv|] where (du, dv) is the gradient
/// rotated into the keypoint frame. Unit normalized.
///
/// Layout: index = (j * 4 + i) * 4 + c, with i along the orientation axis,
/// j perpendicular to it, c the component.
fn describe(img: &GrayF32, x: f32, y: f32, sigma: f32, angle: f32) -> [f32; 64] {
    let (sin, cos) = angle.sin_cos();
    let sig_w = 3.3 * sigma;
    let cell = DESC_SAMPLES / DESC_GRID;
    let mut desc = [0f32; 64];
    for sj in 0..DESC_SAMPLES {
        for si in 0..DESC_SAMPLES {
            let u = (si as f32 - DESC_HALF_EXTENT) * sigma;
            let v = (sj as f32 - DESC_HALF_EXTENT) * sigma;
            let px = x + u * cos - v * sin;
            let py = y + u * sin + v * cos;
            let (gx, gy) = img.gradient(px, py);
            let du = cos * gx + sin * gy;
            let dv = -sin * gx + cos * gy;
            let w = (-(u * u + v * v) / (2.0 * sig_w * sig_w)).exp();
            let base = ((sj / cell) * DESC_GRID + si / cell) * 4;
            desc[base] += w * du;
            desc[base + 1] += w * dv;
            desc[base + 2] += w * du.abs();
            desc[base + 3] += w * dv.abs();
        }
    }
    let norm: f32 = desc.iter().map(|v| v * v).sum::<f32>().sqrt();
    if norm > 0.0 {
        for v in desc.iter_mut() {
            *v /= norm;
        }
    }
    desc
}

/// The descriptor a horizontally mirrored copy of the patch would produce.
///
/// Under a mirror the dominant orientation maps to pi - theta. Working out
/// the sample grid in that frame: the u axis is preserved, the v axis
/// negates. So subregion (i, j) swaps with (i, 3 - j), the dv sum changes
/// sign, and the other three components are untouched.
pub fn mirror_descriptor(d: &[f32; 64]) -> [f32; 64] {
    let mut out = [0f32; 64];
    for j in 0..4 {
        for i in 0..4 {
            let src = (j * 4 + i) * 4;
            let dst = ((3 - j) * 4 + i) * 4;
            out[dst] = d[src];
            out[dst + 1] = -d[src + 1];
            out[dst + 2] = d[src + 2];
            out[dst + 3] = d[src + 3];
        }
    }
    out
}

/// Pick the mirror representative deterministically, so a patch and its
/// mirror image canonicalize to the same descriptor
fn canonicalize(d: [f32; 64]) -> [f32; 64] {
    let dv_total: f32 = (0..16).map(|s| d[s * 4 + 1]).sum();
    let mirror = if dv_total.abs() > 1e-6 {
        dv_total < 0.0
    } else {
        // fallback: skew of the |dv| mass along the v axis
        let skew: f32 = (0..16).map(|s| ((s / 4) as f32 - 1.5) * d[s * 4 + 3]).sum();
        skew < 0.0
    };
    if mirror {
        mirror_descriptor(&d)
    } else {
        d
    }
}

/// 64 bits from one descriptor: the 32 signed sums contribute their sign,
/// the 32 magnitude sums contribute a comparison against their own median.
/// Same recipe as the DCT hash: signs plus ordinal facts, nothing absolute.
pub fn binarize64(desc: &[f32; 64]) -> u64 {
    let mut mags: Vec<f32> = (0..16)
        .flat_map(|s| [desc[s * 4 + 2], desc[s * 4 + 3]])
        .collect();
    mags.sort_by(f32::total_cmp);
    let median = mags[16];
    let mut mask = 0u64;
    for (i, &v) in desc.iter().enumerate() {
        mask <<= 1;
        let bit = if i % 4 < 2 { v > 0.0 } else { v > median };
        if bit {
            mask |= 1;
        }
    }
    mask
}

/// Response-weighted mean of all descriptors, unit normalized.
/// Symmetric pooling, so it inherits every per-keypoint invariance.
pub fn pooled_vector(feats: &[Feature]) -> [f32; 64] {
    let mut acc = [0f32; 64];
    for f in feats {
        for (a, d) in acc.iter_mut().zip(&f.desc) {
            *a += f.kp.response * d;
        }
    }
    let norm: f32 = acc.iter().map(|v| v * v).sum::<f32>().sqrt();
    if norm > 0.0 {
        for v in acc.iter_mut() {
            *v /= norm;
        }
    }
    acc
}

/// Variant 1: binarized pooled descriptor
pub fn hash_pooled(feats: &[Feature]) -> u64 {
    binarize64(&pooled_vector(feats))
}

/// Variant 2: binarized descriptor of the single strongest keypoint
pub fn hash_strongest(feats: &[Feature]) -> u64 {
    match feats.first() {
        Some(f) => binarize64(&f.desc),
        None => 0,
    }
}

/// Variant 3: per-bit majority vote over all binarized descriptors
pub fn hash_consensus(feats: &[Feature]) -> u64 {
    if feats.is_empty() {
        return 0;
    }
    let hashes: Vec<u64> = feats.iter().map(|f| binarize64(&f.desc)).collect();
    let mut mask = 0u64;
    for bit in (0..64).rev() {
        let votes = hashes.iter().filter(|h| (*h >> bit) & 1 == 1).count();
        mask <<= 1;
        if votes * 2 > hashes.len() {
            mask |= 1;
        }
    }
    mask
}

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

    fn dist(a: u64, b: u64) -> u32 {
        (a ^ b).count_ones()
    }

    #[test]
    fn mirror_is_an_involution() {
        let mut d = [0f32; 64];
        for (i, v) in d.iter_mut().enumerate() {
            *v = (i as f32 * 0.37).sin();
        }
        assert_eq!(mirror_descriptor(&mirror_descriptor(&d)), d);
    }

    #[test]
    fn canonical_ignores_mirroring() {
        let mut d = [0f32; 64];
        for (i, v) in d.iter_mut().enumerate() {
            *v = (i as f32 * 0.37).sin();
        }
        assert_eq!(canonicalize(d), canonicalize(mirror_descriptor(&d)));
    }

    #[test]
    fn binarize_packs_all_bits() {
        // strongest positive signal on the first and last positions
        let mut d = [0f32; 64];
        d[0] = 1.0; // first sign bit, lands on bit 63
        d[63] = 1.0; // a magnitude above the (zero) median, lands on bit 0
        let h = binarize64(&d);
        assert_eq!(h, (1 << 63) | 1);
    }

    #[test]
    fn features_deterministic() {
        // synthetic gradients are too smooth to produce DoG extrema
        let img = image::open("img/meowl.jpg").unwrap();
        let s = Salient::new();
        let a = s.features(&img);
        let b = s.features(&img);
        assert_eq!(a.len(), b.len());
        assert!(!a.is_empty());
        for (fa, fb) in a.iter().zip(&b) {
            assert_eq!(fa.desc, fb.desc);
        }
    }

    #[test]
    fn flip_invariant_on_real_photo() {
        let img = image::open("img/meowl.jpg").unwrap();
        let s = Salient::new();
        let a = s.features(&img);
        let b = s.features(&img.fliph());
        assert!(!a.is_empty());
        // detection mirrors exactly, orientation binning wobbles a little
        let d = dist(hash_pooled(&a), hash_pooled(&b));
        assert!(d <= 6, "pooled flip distance {d}");
        let d = dist(hash_consensus(&a), hash_consensus(&b));
        assert!(d <= 6, "consensus flip distance {d}");
    }
}