//! TRACES: FR-CULL-9 | FR-CULL-10
//! Finding the face pairs that could possibly be the same person.
//!
//! [`crate::cluster`] used to begin by computing every pairwise cosine and
//! holding the lot as an `n²` matrix of `f32`. At 1,800 faces that is 13 MB,
//! which is why it survived; at 25,000 it is 2.5 GB, which is why it could not
//! keep surviving.
//!
//! Almost all of that matrix is thrown away unread. Clustering only ever asks
//! whether a pair is *above* the merge threshold, and in a real library the
//! answer is no for well over 99% of pairs — the reference library's 1,813
//! faces produced 7,875 qualifying pairs out of 1.6 million. So this module
//! answers the only question that is actually asked — **which pairs clear the
//! bar** — and returns that sparse list. Memory goes from `O(n²)` to `O(edges)`
//! and the caller never has to hold a matrix at all.
//!
//! # Exact, not approximate
//!
//! The usual way to make this fast is an approximate nearest-neighbour index,
//! which trades recall for speed: it *misses* some true neighbours, and a
//! missed neighbour here is a face that silently never joins its person.
//! Nothing surfaces that — the screen just quietly shows one person as two —
//! so it is a poor trade for a feature whose whole job is to be trusted. This
//! module is exact, and the clusters it produces are identical to those from a
//! full scan.
//!
//! # An exact index was tried, measured, and removed
//!
//! Worth recording so it is not rediscovered as a good idea. The obvious exact
//! index is IVF with a triangle-inequality bound: group the embeddings into
//! cells, and skip a whole cell **pair** when the geometry proves no member of
//! one can reach any member of the other. On the sphere,
//!
//! ```text
//! angle(x, y) >= angle(c_P, c_Q) - radius(P) - radius(Q)
//! ```
//!
//! so a cell pair is impossible when `cos` of that lower bound falls below the
//! threshold. Exact, no recall loss, and it prunes beautifully on synthetic
//! clusters.
//!
//! It prunes **nothing at all** on real face embeddings. Measured over the
//! 1,813-face reference library, at √n = 43 cells:
//!
//! | quantity | measured |
//! |---|---|
//! | median pair angle | 88.5° (cosine 0.026) |
//! | merge threshold | 66.2° (cosine 0.403) |
//! | median cell radius | 80.4° |
//! | median centroid separation | 85.0° |
//! | cell pairs surviving the bound | **946 of 946 — 100%** |
//!
//! The arithmetic is not close. For the bound to exclude a typical cell pair it
//! needs `radius(P) + radius(Q) < 85° - 66° = 19°`, so cells of radius under
//! ~10°. But two photographs of the *same person* sit 36–60° apart, so even a
//! perfect single-identity cell has a radius three times too large. No
//! ball-based partition of this space can have cells tight enough for the
//! inequality to bite — 512-d embeddings are near-orthogonal, and that is the
//! curse of dimensionality doing exactly what it says.
//!
//! So the scan stayed exhaustive, and the effort went where it actually pays:
//! not materialising the matrix, an unrolled dot product, and spreading the
//! blocks across cores. That is `O(n²)` time and `O(edges)` memory, which for
//! this problem is the honest answer.

use crate::calibrate::Calibration;

/// Rows of the similarity triangle handed to one thread at a time.
///
/// Small enough that the tail of the triangle divides evenly across cores —
/// row `i` does `n - i` comparisons, so equal *row counts* are very unequal
/// work — and large enough that the per-block overhead disappears.
const BLOCK: usize = 64;

/// Below this many faces, do the whole thing on the calling thread.
///
/// Spawning threads for a set this small costs more than the scan.
const THREADS_ABOVE: usize = 2048;

/// One face pair that clears the merge threshold.
///
/// `i < j` always, and the probability is carried because the caller would
/// otherwise recompute the sigmoid it took a dot product to reach.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Pair {
    pub i: usize,
    pub j: usize,
    pub probability: f32,
}

/// What the caller has to tell us about each face.
///
/// Deliberately not [`crate::cluster::Candidate`]: this module has no business
/// knowing what a person or a photograph is, and taking the three arrays it
/// actually reads keeps it testable on bare vectors.
pub struct Faces<'a> {
    /// L2-normalised, `EMBEDDING_DIM` long, one per face.
    pub embeddings: &'a [Vec<f32>],
    /// Source pixels across the aligned crop, for the calibration's size term.
    pub crop_px: &'a [f32],
    /// Which photograph each face came from. Two faces in one frame are not
    /// the same person, so those pairs are never returned (docs/faces.md §9).
    pub images: &'a [u64],
}

/// Every pair whose calibrated probability reaches `min_probability`.
///
/// Excludes pairs from the same photograph, which the clusterer would refuse
/// anyway — dropping them here keeps them out of the graph the caller builds
/// and out of the connected components it derives from it.
///
/// Ordered by `(i, j)`, which is what the caller's determinism rests on.
pub fn above_threshold(faces: &Faces, cal: &Calibration, min_probability: f32) -> Vec<Pair> {
    let n = faces.embeddings.len();
    if n < 2 {
        return Vec::new();
    }

    // The loosest cosine that could clear the bar for *any* pair in the set.
    // Cheaper than the sigmoid by far, and it rejects almost everything.
    let tau = loosest_cosine(faces.crop_px, cal, min_probability);

    let blocks: Vec<(usize, usize)> = (0..n)
        .step_by(BLOCK)
        .map(|start| (start, (start + BLOCK).min(n)))
        .collect();

    if n < THREADS_ABOVE {
        let mut out = Vec::new();
        for &(from, to) in &blocks {
            scan_block(faces, cal, min_probability, tau, from, to, &mut out);
        }
        return out;
    }

    // One worker per core bar one. This runs on a background thread behind a
    // button the user pressed, and NFR-ARCH-2 puts it behind the UI: taking
    // every core would stall the window it is reporting progress to.
    let workers = std::thread::available_parallelism()
        .map(|p| p.get().saturating_sub(1).max(1))
        .unwrap_or(1)
        .min(blocks.len());

    let next = std::sync::atomic::AtomicUsize::new(0);
    let mut parts: Vec<Vec<Vec<Pair>>> = std::thread::scope(|scope| {
        let handles: Vec<_> = (0..workers)
            .map(|_| {
                let next = &next;
                let blocks = &blocks;
                scope.spawn(move || {
                    // Results stay tagged with their block index, so the order
                    // of the output does not depend on which thread got there
                    // first. Determinism is a promise this module keeps.
                    let mut mine: Vec<(usize, Vec<Pair>)> = Vec::new();
                    loop {
                        let b = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                        let Some(&(from, to)) = blocks.get(b) else {
                            break;
                        };
                        let mut out = Vec::new();
                        scan_block(faces, cal, min_probability, tau, from, to, &mut out);
                        mine.push((b, out));
                    }
                    mine
                })
            })
            .collect();

        let mut slots: Vec<Vec<Vec<Pair>>> = vec![Vec::new(); blocks.len()];
        for h in handles {
            for (b, pairs) in h.join().unwrap_or_default() {
                slots[b].push(pairs);
            }
        }
        slots
    });

    let mut out = Vec::new();
    for slot in &mut parts {
        for pairs in slot.drain(..) {
            out.extend(pairs);
        }
    }
    out
}

/// Compare rows `from..to` against everything after them.
///
/// The upper triangle, split by rows. Row `i` only looks at `j > i`, so every
/// unordered pair is visited exactly once and the emitted order is `(i, j)`
/// ascending within the block.
fn scan_block(
    faces: &Faces,
    cal: &Calibration,
    min_probability: f32,
    tau: f32,
    from: usize,
    to: usize,
    out: &mut Vec<Pair>,
) {
    let n = faces.embeddings.len();
    for i in from..to {
        let a = &faces.embeddings[i];
        let crop_a = faces.crop_px[i];
        let image_a = faces.images[i];
        for j in i + 1..n {
            if image_a == faces.images[j] {
                continue;
            }
            let cos = dot(a, &faces.embeddings[j]);
            // The cheap rejection, and it takes well over 99% of pairs.
            if cos < tau {
                continue;
            }
            let probability = cal.probability(cos, crop_a.min(faces.crop_px[j]), 0.0);
            if probability >= min_probability {
                out.push(Pair { i, j, probability });
            }
        }
    }
}

/// The lowest cosine that could yield `min_probability` for any pair in the set.
///
/// The calibration is `sigmoid(a·cos + b + w_size·log2(crop))`, so for a fixed
/// size term the cosine boundary is exact. The size term is *not* fixed — it
/// varies per pair with the smaller of the two faces — so the safe bound uses
/// whichever face size pushes the boundary lowest: the largest face when
/// `w_size` is positive, the smallest when it is negative.
///
/// Returns [`f32::NEG_INFINITY`] — a filter that rejects nothing — where no
/// boundary exists: a non-positive steepness, for which probability does not
/// increase with cosine, or a threshold at the ends of the sigmoid. Those are
/// degenerate calibrations rather than impossible ones, and the right response
/// is to stop pruning, not to guess.
fn loosest_cosine(crop_px: &[f32], cal: &Calibration, min_probability: f32) -> f32 {
    if cal.a <= 0.0 || !(min_probability > 0.0 && min_probability < 1.0) {
        return f32::NEG_INFINITY;
    }

    let (mut lo, mut hi) = (f32::INFINITY, 0.0_f32);
    for &c in crop_px {
        let c = c.max(1.0);
        lo = lo.min(c);
        hi = hi.max(c);
    }
    if !lo.is_finite() {
        return f32::NEG_INFINITY;
    }

    let extreme = if cal.w_size >= 0.0 { hi } else { lo };
    let tau = cal.boundary_at(min_probability, extreme, 0.0);
    if tau.is_nan() {
        return f32::NEG_INFINITY;
    }
    // Cosines never exceed 1, so a boundary above it legitimately rejects
    // everything. Clamped rather than left free so the comparison stays cheap.
    tau.min(1.0)
}

/// Cosine of two L2-normalised embeddings.
///
/// Eight accumulators rather than one. Floating-point addition is not
/// associative, so the compiler may not re-associate a single running total and
/// the loop serialises on the adder's latency; eight independent chains give it
/// something to pipeline and vectorise. The order is fixed and identical on
/// every run, which is what the caller's determinism needs — it is a different
/// order from the naive sum, not a variable one.
pub(crate) fn dot(a: &[f32], b: &[f32]) -> f32 {
    const LANES: usize = 8;
    let mut acc = [0.0_f32; LANES];
    let chunks = a.len() / LANES;

    for c in 0..chunks {
        let base = c * LANES;
        for (l, slot) in acc.iter_mut().enumerate() {
            *slot += a[base + l] * b[base + l];
        }
    }
    let mut total =
        ((acc[0] + acc[1]) + (acc[2] + acc[3])) + ((acc[4] + acc[5]) + (acc[6] + acc[7]));
    for k in chunks * LANES..a.len() {
        total += a[k] * b[k];
    }
    total
}

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

    const DIM: usize = 64;

    fn cal() -> Calibration {
        Calibration {
            a: 30.0,
            b: -30.0 * 0.35,
            w_size: 0.0,
            valid: true,
            positive_pairs: 1000,
            negative_pairs: 10_000,
        }
    }

    /// A unit vector in a reproducible pseudo-random direction.
    ///
    /// Hashed from the seed rather than drawn from an RNG, so a failure is
    /// reproducible from the test alone.
    fn vector(seed: u64) -> Vec<f32> {
        let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
        let mut v = Vec::with_capacity(DIM);
        for _ in 0..DIM {
            s ^= s << 13;
            s ^= s >> 7;
            s ^= s << 17;
            v.push(((s >> 11) as f64 / (1u64 << 53) as f64) as f32 - 0.5);
        }
        normalise(v)
    }

    fn normalise(mut v: Vec<f32>) -> Vec<f32> {
        let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        for x in &mut v {
            *x /= n;
        }
        v
    }

    /// A vector a known cosine away from `base`.
    fn near(base: &[f32], other: &[f32], cosine: f32) -> Vec<f32> {
        let d = dot(base, other);
        let mut perp: Vec<f32> = other.iter().zip(base).map(|(o, b)| o - d * b).collect();
        let n = perp.iter().map(|x| x * x).sum::<f32>().sqrt();
        for x in &mut perp {
            *x /= n;
        }
        let s = (1.0 - cosine * cosine).max(0.0).sqrt();
        normalise(
            base.iter()
                .zip(&perp)
                .map(|(b, p)| cosine * b + s * p)
                .collect(),
        )
    }

    struct Set {
        embeddings: Vec<Vec<f32>>,
        crop_px: Vec<f32>,
        images: Vec<u64>,
    }

    impl Set {
        fn faces(&self) -> Faces<'_> {
            Faces {
                embeddings: &self.embeddings,
                crop_px: &self.crop_px,
                images: &self.images,
            }
        }
    }

    /// `groups` identities, `per` faces each, every face in its own photograph.
    fn population(groups: usize, per: usize, tightness: f32) -> Set {
        let mut embeddings = Vec::new();
        let mut images = Vec::new();
        let mut image = 0u64;
        for g in 0..groups {
            let base = vector(g as u64 + 1);
            let off = vector(g as u64 + 9_999);
            for m in 0..per {
                embeddings.push(if m == 0 {
                    base.clone()
                } else {
                    near(&base, &off, tightness)
                });
                images.push(image);
                image += 1;
            }
        }
        let crop_px = vec![150.0; embeddings.len()];
        Set {
            embeddings,
            crop_px,
            images,
        }
    }

    /// The unpruned, unthreaded, unblocked definition of the answer.
    fn reference(faces: &Faces, cal: &Calibration, min_probability: f32) -> Vec<Pair> {
        let n = faces.embeddings.len();
        let mut out = Vec::new();
        for i in 0..n {
            for j in i + 1..n {
                if faces.images[i] == faces.images[j] {
                    continue;
                }
                let cos: f32 = faces.embeddings[i]
                    .iter()
                    .zip(&faces.embeddings[j])
                    .map(|(x, y)| x * y)
                    .sum();
                let p = cal.probability(cos, faces.crop_px[i].min(faces.crop_px[j]), 0.0);
                if p >= min_probability {
                    out.push(Pair {
                        i,
                        j,
                        probability: p,
                    });
                }
            }
        }
        out
    }

    fn same_pairs(a: &[Pair], b: &[Pair]) -> bool {
        a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.i == y.i && x.j == y.j)
    }

    #[test]
    fn nothing_to_pair_is_no_pairs() {
        let s = population(1, 1, 0.9);
        assert!(above_threshold(&s.faces(), &cal(), 0.9).is_empty());
    }

    #[test]
    fn the_blocked_scan_finds_exactly_what_the_reference_does() {
        let s = population(60, 8, 0.97);
        let f = s.faces();
        let got = above_threshold(&f, &cal(), 0.9);
        let want = reference(&f, &cal(), 0.9);
        assert!(!want.is_empty(), "the reference found nothing to check");
        assert!(same_pairs(&got, &want), "{} vs {}", got.len(), want.len());
    }

    /// Past `THREADS_ABOVE` the work is split across cores and stitched back
    /// together, and the stitching is where an order bug would live.
    #[test]
    fn the_threaded_scan_finds_exactly_what_the_reference_does() {
        let s = population(300, 8, 0.97);
        assert!(
            s.embeddings.len() > THREADS_ABOVE,
            "population is below the threading cutoff"
        );
        let f = s.faces();
        let got = above_threshold(&f, &cal(), 0.9);
        let want = reference(&f, &cal(), 0.9);
        assert!(!want.is_empty());
        assert!(same_pairs(&got, &want), "{} vs {}", got.len(), want.len());
    }

    /// The size term moves the cosine boundary per pair, so the pre-filter has
    /// to be built from the most permissive size in the set or it will drop a
    /// pair that would have qualified.
    #[test]
    fn a_size_weighted_calibration_still_matches_the_reference() {
        let mut s = population(60, 8, 0.97);
        for (i, c) in s.crop_px.iter_mut().enumerate() {
            *c = 40.0 + (i % 17) as f32 * 30.0;
        }
        let sized = Calibration {
            w_size: 0.5,
            b: -30.0 * 0.35 - 0.5 * 7.0,
            ..cal()
        };
        let f = s.faces();
        assert!(same_pairs(
            &above_threshold(&f, &sized, 0.9),
            &reference(&f, &sized, 0.9)
        ));
    }

    /// A negative size weight flips which extreme is permissive. Cheap to get
    /// wrong and silent when it is, so it gets its own case.
    #[test]
    fn a_negative_size_weight_prunes_from_the_other_end() {
        let mut s = population(60, 8, 0.97);
        for (i, c) in s.crop_px.iter_mut().enumerate() {
            *c = 40.0 + (i % 17) as f32 * 30.0;
        }
        let sized = Calibration {
            w_size: -0.5,
            b: -30.0 * 0.35 + 0.5 * 7.0,
            ..cal()
        };
        let f = s.faces();
        assert!(same_pairs(
            &above_threshold(&f, &sized, 0.9),
            &reference(&f, &sized, 0.9)
        ));
    }

    /// A degenerate calibration has no cosine boundary to prune against, and
    /// must stop pruning rather than prune on a bound that does not hold.
    #[test]
    fn a_flat_calibration_prunes_nothing_and_still_agrees() {
        let flat = Calibration {
            a: 0.0,
            b: 4.0,
            ..cal()
        };
        assert_eq!(loosest_cosine(&[150.0], &flat, 0.9), f32::NEG_INFINITY);
        let s = population(20, 6, 0.97);
        let f = s.faces();
        assert!(same_pairs(
            &above_threshold(&f, &flat, 0.9),
            &reference(&f, &flat, 0.9)
        ));
    }

    #[test]
    fn two_faces_in_one_photograph_are_never_paired() {
        let mut s = population(1, 2, 1.0);
        s.images = vec![7, 7];
        assert!(above_threshold(&s.faces(), &cal(), 0.9).is_empty());
    }

    #[test]
    fn pairs_come_back_in_index_order() {
        let s = population(300, 8, 0.97);
        let pairs = above_threshold(&s.faces(), &cal(), 0.9);
        assert!(pairs
            .windows(2)
            .all(|w| (w[0].i, w[0].j) < (w[1].i, w[1].j)));
        assert!(pairs.iter().all(|p| p.i < p.j));
    }

    /// Threads must not make the answer depend on which one finished first.
    #[test]
    fn the_same_input_yields_the_same_pairs() {
        let s = population(300, 8, 0.97);
        assert_eq!(
            above_threshold(&s.faces(), &cal(), 0.9),
            above_threshold(&s.faces(), &cal(), 0.9)
        );
    }

    /// The unrolled dot has to agree with the obvious one, tail included — the
    /// lengths here are deliberately not multiples of the lane count.
    #[test]
    fn the_unrolled_dot_matches_the_naive_one() {
        for len in [1usize, 7, 8, 9, 63, 64, 65, 512] {
            let a: Vec<f32> = (0..len).map(|i| (i as f32 * 0.37).sin()).collect();
            let b: Vec<f32> = (0..len).map(|i| (i as f32 * 0.11).cos()).collect();
            let naive: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
            assert!(
                (dot(&a, &b) - naive).abs() < 1e-4,
                "len {len}: {} vs {naive}",
                dot(&a, &b)
            );
        }
    }

    /// The pre-filter is the whole speed story, so it is worth asserting it
    /// actually rejects the bulk of the population rather than trusting it to.
    #[test]
    fn the_threshold_filter_rejects_almost_everything() {
        let s = population(60, 8, 0.97);
        let n = s.embeddings.len();
        let total = n * (n - 1) / 2;
        let kept = above_threshold(&s.faces(), &cal(), 0.9).len();
        assert!(
            kept * 20 < total,
            "kept {kept} of {total} pairs, which is not sparse"
        );
    }
}