//! What an embedder produces, and how it is stored (docs/dev/faces.md §6).
//!
//! Deliberately **model-free**: the vector, its identity, its comparison and
//! its storage encoding are arithmetic, and `calibrate` and `cluster` are built
//! on them. Keeping them out of the `inference` feature is what lets the part
//! of this subsystem most likely to be subtly wrong be tested on a machine with
//! no weights on it.
//!
//! [`crate::embed::Embedder`] is the thing that needs a model, and it lives
//! behind the feature.

/// Embedding dimensionality. Fixed by the model family, not a parameter.
pub const EMBEDDING_DIM: usize = 512;

/// The shortest raw embedding a face may be *compared against*.
///
/// # What the length of the vector says
///
/// ArcFace is trained on the direction of its output and nothing else, and
/// the length it leaves behind turns out to be a free quality signal: the
/// magnitude grows with how recognisable the crop was to the model, and a
/// blurred, occluded, badly lit or hard-profile face comes out short. MagFace
/// (Meng et al., CVPR 2021) made that the training objective; the plain
/// ArcFace heads this crate runs already show it, weaker but usable, which is
/// why it is worth keeping the number the normalisation discards.
///
/// # Why it gates the gallery and not the face
///
/// A short vector is a bad *reference*: it sits nearer the centre of the
/// sphere than a real identity does and matches a little of everyone, which
/// is exactly the face that welds two people together in a clustering pass.
/// It is not a bad *probe* — the face is still real, still somebody, and
/// comparing it against good references is the only way it will ever be named.
/// So a face below this floor is compared against the gallery and never
/// becomes part of it: see `cluster::Candidate::in_gallery`.
///
/// 14 is the operating point for `w600k_mbf`, whose norms on the reference
/// library run from about 8 on a blur to the high 20s on a clean portrait. A
/// face whose quality was never recorded — indexed before the number was kept
/// — is not gated, because a rule that cannot be checked should admit, not
/// exclude.
pub const MIN_GALLERY_QUALITY: f32 = 14.0;

/// Whether an embedding of this quality may serve as a reference.
///
/// `None` is "not measured", and is admitted: the rule is about a number that
/// was read and found short, not about a number that is missing.
pub fn in_gallery(quality: Option<f32>) -> bool {
    quality.is_none_or(|q| q >= MIN_GALLERY_QUALITY)
}

/// Which model produced an embedding.
///
/// Embeddings from different models are not comparable, and this is the one
/// mistake that produces plausible-looking garbage rather than an error — so
/// the id travels *with* the vector rather than beside it, and
/// [`Embedding::cosine`] refuses a cross-model comparison.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ModelId(pub std::sync::Arc<str>);

impl ModelId {
    pub fn new(s: impl Into<std::sync::Arc<str>>) -> Self {
        Self(s.into())
    }
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for ModelId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

/// A 512-d L2-normalised face embedding.
#[derive(Debug, Clone, PartialEq)]
pub struct Embedding {
    pub model: ModelId,
    pub v: Box<[f32; EMBEDDING_DIM]>,
}

impl Embedding {
    /// Cosine similarity, which for unit vectors is the plain dot product.
    ///
    /// `None` when the two came from different models. That is a real
    /// possibility in a library indexed across a model upgrade, and the
    /// alternative — returning a number — is the failure mode
    /// `faces.model_id` exists to prevent.
    pub fn cosine(&self, other: &Embedding) -> Option<f32> {
        if self.model != other.model {
            return None;
        }
        Some(dot(&self.v, &other.v))
    }

    /// Storage form: `512 × f16`, 1 KB per face (catalog.md §10.1).
    ///
    /// This writes the unit vector. What the catalog stores is the raw one —
    /// `embed::Embedded::to_f16_bytes` — because the length is the quality
    /// and a unit vector has none left to read.
    pub fn to_f16_bytes(&self) -> Vec<u8> {
        self.to_f16_bytes_scaled(1.0)
    }

    /// The unit vector scaled by `length`, as `512 × f16`.
    pub(crate) fn to_f16_bytes_scaled(&self, length: f32) -> Vec<u8> {
        let mut out = Vec::with_capacity(EMBEDDING_DIM * 2);
        for &x in self.v.iter() {
            out.extend_from_slice(&f32_to_f16_bits(x * length).to_le_bytes());
        }
        out
    }

    /// Read back from storage, re-normalising.
    ///
    /// The f16 round-trip perturbs a unit vector by ~1e-3 in cosine — three
    /// orders below the separation between a match and a non-match — but the
    /// drift is free to remove and invisible if left, so it is removed here
    /// rather than remembered at every call site. The same pass is what turns
    /// a stored raw vector back into the unit one every comparison expects.
    pub fn from_f16_bytes(model: ModelId, bytes: &[u8]) -> Option<Self> {
        read_f16_bytes(model, bytes).map(|(e, _)| e)
    }
}

/// Read a stored vector back, with the length it was stored at.
///
/// The length is the quality where the blob is a raw one, and ~1 where it is
/// a unit vector from before raw vectors were stored — which is why the
/// catalog keeps the quality beside the blob rather than deriving it from
/// this: a unit vector reads as a quality of 1, not as "unmeasured".
pub fn read_f16_bytes(model: ModelId, bytes: &[u8]) -> Option<(Embedding, f32)> {
    if bytes.len() != EMBEDDING_DIM * 2 {
        return None;
    }
    let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
    for (i, chunk) in bytes.chunks_exact(2).enumerate() {
        v[i] = f16_bits_to_f32(u16::from_le_bytes([chunk[0], chunk[1]]));
    }
    let length = normalise(&mut v);
    Some((Embedding { model, v }, length))
}

fn dot(a: &[f32; EMBEDDING_DIM], b: &[f32; EMBEDDING_DIM]) -> f32 {
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

/// Scale `v` to unit length, and return the length it had.
///
/// The length is the one thing about the raw output that survives being
/// thrown away by everything downstream, and it is a quality signal
/// ([`MIN_GALLERY_QUALITY`]) — so it comes back out rather than being lost
/// here.
pub(crate) fn normalise(v: &mut [f32; EMBEDDING_DIM]) -> f32 {
    // Clamped rather than checked: a zero-norm embedding is a broken model,
    // not a runtime condition worth an error path, and dividing by 1e-6 keeps
    // the NaN out of the catalog.
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-6);
    for x in v.iter_mut() {
        *x /= norm;
    }
    norm
}

// ── f16 ───────────────────────────────────────────────────────────────────
//
// Hand-rolled rather than pulling in `half`: two functions over a format that
// has not changed since 2008, used at exactly one boundary. The dependency
// policy (D13, D1) makes the bar for a new crate high, and this is well under
// it.

fn f32_to_f16_bits(x: f32) -> u16 {
    let bits = x.to_bits();
    let sign = ((bits >> 16) & 0x8000) as u16;
    let exp = ((bits >> 23) & 0xff) as i32 - 127 + 15;
    let mant = bits & 0x007f_ffff;

    if exp >= 0x1f {
        // Overflow, inf, or NaN. No component of an embedding exceeds its
        // length, and the lengths this model produces are in the tens, so
        // this is the broken-model path; infinity is the honest answer, not a
        // clamp that hides it.
        return sign
            | 0x7c00
            | if mant != 0 && exp == 0x1f + 112 {
                0x200
            } else {
                0
            };
    }
    if exp <= 0 {
        // Subnormal or underflow. A component of a unit 512-vector is ~0.04
        // and a stored one is that times the length, nowhere near here, so
        // this branch exists for correctness rather than for traffic.
        if exp < -10 {
            return sign;
        }
        let mant = mant | 0x0080_0000;
        let shift = (14 - exp) as u32;
        let half = (mant >> shift) as u16;
        // Round to nearest, ties to even.
        let rem = mant & ((1 << shift) - 1);
        let tie = 1 << (shift - 1);
        let round = u16::from(rem > tie || (rem == tie && (half & 1) == 1));
        return sign | (half + round);
    }

    let half = ((exp as u16) << 10) | (mant >> 13) as u16;
    let rem = mant & 0x1fff;
    let round = u16::from(rem > 0x1000 || (rem == 0x1000 && (half & 1) == 1));
    sign | (half + round)
}

fn f16_bits_to_f32(h: u16) -> f32 {
    let sign = ((h & 0x8000) as u32) << 16;
    let exp = ((h >> 10) & 0x1f) as u32;
    let mant = (h & 0x03ff) as u32;

    if exp == 0 {
        if mant == 0 {
            return f32::from_bits(sign);
        }
        // Subnormal: renormalise into f32's range.
        let mut e = -1_i32;
        let mut m = mant;
        while m & 0x0400 == 0 {
            m <<= 1;
            e -= 1;
        }
        let m = m & 0x03ff;
        return f32::from_bits(sign | (((127 - 15 + 1 + e) as u32) << 23) | (m << 13));
    }
    if exp == 0x1f {
        return f32::from_bits(sign | 0x7f80_0000 | (mant << 13));
    }
    f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mant << 13))
}

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

    fn unit(seed: u32) -> Embedding {
        let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
        let mut s = seed.wrapping_mul(2_654_435_761).wrapping_add(1);
        for x in v.iter_mut() {
            s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            *x = (s >> 8) as f32 / (1u32 << 23) as f32 - 0.5;
        }
        normalise(&mut v);
        Embedding {
            model: ModelId::new("test"),
            v,
        }
    }

    #[test]
    fn a_normalised_embedding_has_cosine_one_with_itself() {
        let e = unit(7);
        assert!((e.cosine(&e).unwrap() - 1.0).abs() < 1e-5);
    }

    #[test]
    fn embeddings_from_different_models_do_not_compare() {
        let a = unit(1);
        let mut b = unit(1);
        b.model = ModelId::new("other");
        assert_eq!(
            a.cosine(&b),
            None,
            "a cross-model cosine must not be a number"
        );
    }

    /// The claim docs/dev/faces.md §6 makes about the storage format: the f16
    /// round-trip costs ~1e-3 of cosine, three orders below the separation
    /// between a match and a non-match.
    #[test]
    fn f16_round_trip_preserves_the_embedding() {
        for seed in 0..16 {
            let e = unit(seed);
            let back = Embedding::from_f16_bytes(e.model.clone(), &e.to_f16_bytes()).unwrap();
            let cos = e.cosine(&back).unwrap();
            assert!(cos > 0.9999, "seed {seed}: round-trip cosine {cos}");
        }
    }

    #[test]
    fn normalising_reports_the_length_it_removed() {
        let mut v = Box::new([0.0_f32; EMBEDDING_DIM]);
        v[0] = 3.0;
        v[1] = 4.0;
        let norm = normalise(&mut v);
        assert!((norm - 5.0).abs() < 1e-6, "norm {norm}");
        assert!((v[0] - 0.6).abs() < 1e-6 && (v[1] - 0.8).abs() < 1e-6);
    }

    /// The gate admits what it cannot measure: a face from before the number
    /// was kept is not a face that was found wanting.
    #[test]
    fn an_unmeasured_quality_is_admitted_to_the_gallery() {
        assert!(in_gallery(None));
        assert!(in_gallery(Some(MIN_GALLERY_QUALITY)));
        assert!(in_gallery(Some(27.5)));
        assert!(!in_gallery(Some(MIN_GALLERY_QUALITY - 0.01)));
        assert!(!in_gallery(Some(8.0)));
    }

    /// The storage form carries the length, and the length comes back out —
    /// without touching the direction every comparison is made on.
    #[test]
    fn a_raw_vector_round_trips_with_its_length() {
        let e = unit(3);
        let raw = e.to_f16_bytes_scaled(21.5);
        let (back, length) = read_f16_bytes(e.model.clone(), &raw).unwrap();
        assert!((length - 21.5).abs() < 0.05, "length {length}");
        assert!(e.cosine(&back).unwrap() > 0.9999);
        // A unit vector from an older store reads as length 1, not as an
        // error — see `read_f16_bytes` on why that is not "unmeasured".
        let (_, one) = read_f16_bytes(e.model.clone(), &e.to_f16_bytes()).unwrap();
        assert!((one - 1.0).abs() < 1e-2, "length {one}");
    }

    #[test]
    fn f16_round_trip_rejects_a_wrong_length_blob() {
        assert!(Embedding::from_f16_bytes(ModelId::new("m"), &[0u8; 100]).is_none());
    }

    #[test]
    fn f16_handles_the_values_an_embedding_actually_contains() {
        // Components of a unit 512-vector cluster around ±1/sqrt(512) ≈ 0.044.
        for &x in &[0.0_f32, 1.0, -1.0, 0.044_194_17, -0.044_194_17, 1e-3, -7e-4] {
            let back = f16_bits_to_f32(f32_to_f16_bits(x));
            assert!(
                (back - x).abs() <= 1e-3 * x.abs().max(1e-3),
                "{x} round-tripped to {back}"
            );
        }
    }
}