//! Preference vector: per-user taste embedding with L2 normalization invariant.
//!
//! Tracks user taste by maintaining a preference vector that evolves with
//! interactions. The vector is L2-normalized on every update to ensure
//! consistent cosine similarity scoring during personalized ranking.
//!
//! # Durability
//!
//! Preference vectors are DERIVED state with no WAL backing (the signal record
//! carries no `for_user`), so they are persisted by [`PreferenceVectors::checkpoint`]
//! (called on the periodic checkpoint thread, on backup, and on shutdown) and
//! restored on open by [`PreferenceVectors::restore`]. This bounds the post-crash
//! loss window to the checkpoint interval rather than losing every update since
//! the last clean shutdown.

use dashmap::DashMap;

use crate::storage::vector::{VectorError, l2_normalize_in_place};

/// Per-user preference vector, L2-normalized.
///
/// The vector is updated via exponential moving average: each new interaction
/// embedding is blended with the current preference using a learning rate.
///
/// Thread-safe via `DashMap` -- concurrent updates to different users never
/// contend.
pub struct PreferenceVectors {
    /// `user_id` -> normalized preference vector
    inner: DashMap<u64, Vec<f32>>,
    /// Dimensionality of the embedding space. All vectors must have this length.
    dim: usize,
    /// Base learning rate for exponential moving average updates.
    /// Default: 0.1. The adaptive rate decays as `base / (1 + ln(count + 1))`.
    learning_rate: f32,
    /// Per-user update counts for adaptive learning rate computation.
    ///
    /// # Known Limitation
    ///
    /// `update_counts` is in-memory only and is NOT persisted to storage. On
    /// process restart, all users start at count=0 and their adaptive learning
    /// rate resets to `base_alpha`. This means the learning rate decays more
    /// slowly across session boundaries than within a single session.
    /// Persisting update counts requires a storage write on every preference
    /// update, which introduces non-trivial overhead on the hot path; this
    /// trade-off is deferred to M7 (Production Hardening).
    update_counts: DashMap<u64, u64>,
}

impl PreferenceVectors {
    /// Create a new preference vector store for the given embedding dimensionality.
    #[must_use]
    pub fn new(dim: usize) -> Self {
        Self {
            inner: DashMap::new(),
            dim,
            learning_rate: 0.1,
            update_counts: DashMap::new(),
        }
    }

    /// Create with a custom base learning rate.
    #[must_use]
    pub fn with_learning_rate(dim: usize, learning_rate: f32) -> Self {
        Self {
            inner: DashMap::new(),
            dim,
            learning_rate,
            update_counts: DashMap::new(),
        }
    }

    /// Get the current preference vector for a user (cloned).
    ///
    /// Returns `None` if no preference has been recorded.
    #[must_use]
    pub fn get(&self, user_id: u64) -> Option<Vec<f32>> {
        self.inner.get(&user_id).map(|r| r.clone())
    }

    /// Set the preference vector directly (e.g., from cold-start initialization).
    ///
    /// The vector is L2-normalized before storage. Returns `false` if the
    /// dimension does not match.
    #[must_use]
    pub fn set(&self, user_id: u64, mut vec: Vec<f32>) -> bool {
        if vec.len() != self.dim {
            return false;
        }
        normalize_centroid(&mut vec);
        self.inner.insert(user_id, vec);
        true
    }

    /// Update a user's preference vector by blending with an interaction embedding.
    ///
    /// Uses exponential moving average with an adaptive learning rate:
    /// `alpha = base_alpha / (1 + ln(update_count + 1))`
    /// `pref = (1 - alpha) * pref + alpha * interaction`
    /// then L2-normalizes the result.
    ///
    /// The adaptive rate decays logarithmically with the number of updates,
    /// so early interactions have high influence and later interactions refine
    /// the preference more gently.
    ///
    /// If no preference exists yet, the interaction embedding becomes the initial
    /// preference (after normalization). Uses `Entry::Occupied`/`Entry::Vacant`
    /// to avoid double-applying the blend on first insertion.
    ///
    /// Returns `false` if the interaction embedding dimension does not match.
    #[must_use]
    pub fn update(&self, user_id: u64, interaction_embedding: &[f32]) -> bool {
        if interaction_embedding.len() != self.dim {
            return false;
        }

        // Compute adaptive learning rate: alpha = base / (1 + ln(count + 1)).
        let count = {
            let mut entry = self.update_counts.entry(user_id).or_insert(0);
            let c = *entry;
            *entry += 1;
            c
        };
        #[allow(clippy::cast_precision_loss)]
        let lr = (f64::from(self.learning_rate) / (1.0 + (count as f64).ln_1p())) as f32;

        self.blend_and_normalize(user_id, interaction_embedding, lr);
        true
    }

    /// Update a user's preference vector by blending with an interaction embedding
    /// using an explicit learning rate (overrides the stored rate).
    ///
    /// Same EMA formula as [`update`](Self::update) but with `lr` parameter:
    /// `pref = (1 - lr) * pref + lr * interaction`
    ///
    /// Returns `false` if the interaction embedding dimension does not match.
    #[must_use]
    pub fn update_with_custom_rate(
        &self,
        user_id: u64,
        interaction_embedding: &[f32],
        lr: f32,
    ) -> bool {
        if interaction_embedding.len() != self.dim {
            return false;
        }

        self.blend_and_normalize(user_id, interaction_embedding, lr);
        true
    }

    /// Blend `interaction_embedding` into the user's preference at learning
    /// rate `lr` and re-normalize, the shared EMA body for both [`update`] and
    /// [`update_with_custom_rate`].
    ///
    /// An existing preference is blended in place via
    /// `pref = (1 - lr) * pref + lr * interaction` and re-normalized; if no
    /// preference exists yet, the (normalized) interaction embedding becomes the
    /// initial preference, avoiding a double-applied blend on first insertion.
    ///
    /// The caller is responsible for the dimension check — `dim` mismatch is
    /// validated before this is reached, and the per-element `zip` would
    /// otherwise silently blend only the overlapping prefix.
    ///
    /// [`update`]: Self::update
    /// [`update_with_custom_rate`]: Self::update_with_custom_rate
    fn blend_and_normalize(&self, user_id: u64, interaction_embedding: &[f32], lr: f32) {
        use dashmap::mapref::entry::Entry;

        match self.inner.entry(user_id) {
            Entry::Occupied(mut occ) => {
                let pref = occ.get_mut();
                for (p, &i) in pref.iter_mut().zip(interaction_embedding.iter()) {
                    *p = (1.0 - lr).mul_add(*p, lr * i);
                }
                normalize_centroid(pref);
            }
            Entry::Vacant(vac) => {
                let mut v = interaction_embedding.to_vec();
                normalize_centroid(&mut v);
                vac.insert(v);
            }
        }
    }

    /// Compute cosine similarity between a user's preference and a candidate embedding.
    ///
    /// Returns `None` if the user has no preference vector or dimensions mismatch.
    /// The stored preference is L2-normalized; the candidate is normalized on-the-fly
    /// so callers do not need to pre-normalize.
    #[must_use]
    #[allow(clippy::significant_drop_tightening)]
    pub fn cosine_similarity(&self, user_id: u64, candidate: &[f32]) -> Option<f32> {
        if candidate.len() != self.dim {
            return None;
        }
        let pref = self.inner.get(&user_id)?;
        let dot: f32 = pref.iter().zip(candidate.iter()).map(|(a, b)| a * b).sum();
        // Divide by the candidate's L2 norm to get true cosine similarity.
        // The stored preference is already unit-length, so we only need
        // to normalize the candidate side.
        let candidate_norm: f32 = candidate.iter().map(|x| x * x).sum::<f32>().sqrt();
        if candidate_norm < f32::EPSILON {
            return Some(0.0);
        }
        Some(dot / candidate_norm)
    }

    /// Remove every stored preference vector and update count.
    ///
    /// Used to return the store to its just-opened state so a full rebuild
    /// (re-seed + signal replay) does not fold on top of already-folded vectors
    /// — i.e. to keep an out-of-band reindex idempotent with a fresh boot.
    pub fn clear(&self) {
        self.inner.clear();
        self.update_counts.clear();
    }

    /// Remove a single user's cold-start vector and update count.
    ///
    /// Used by [`MultiPreferenceVectors`](crate::entities::MultiPreferenceVectors)
    /// when a user crosses the warm threshold: their cold-start vector is migrated
    /// into cluster 0 and then dropped here so a checkpoint never double-stores a
    /// stale K=1 row alongside the cluster rows.
    pub fn remove(&self, user_id: u64) {
        self.inner.remove(&user_id);
        self.update_counts.remove(&user_id);
    }

    /// Insert an already-restored (normalized) vector + update count directly,
    /// bypassing the EMA blend.
    ///
    /// The multi-vector restore path uses this to load a legacy single-vector
    /// checkpoint row as a cold-start user. The caller is responsible for having
    /// re-normalized the vector at the load boundary (matching [`restore`]'s
    /// contract); the dimension is trusted because the caller already gated on it.
    ///
    /// [`restore`]: Self::restore
    pub fn insert_restored(&self, user_id: u64, vec: Vec<f32>, update_count: u64) {
        self.inner.insert(user_id, vec);
        self.update_counts.insert(user_id, update_count);
    }

    /// Whether this store holds a vector for `user_id`, without allocating
    /// (an O(1) hash lookup). Used by
    /// [`MultiPreferenceVectors::contains`](crate::entities::MultiPreferenceVectors::contains)
    /// for serve-path existence guards.
    #[must_use]
    pub fn contains(&self, user_id: u64) -> bool {
        self.inner.contains_key(&user_id)
    }

    /// Stage one legacy single-vector row per stored user into an existing
    /// `WriteBatch` (`[count:8 LE][dim:4 LE][f32*dim]`), without performing its
    /// own delete sweep. `skip(user_id) == true` omits that user's row.
    ///
    /// [`checkpoint`](Self::checkpoint) is the standalone single-vector path (it
    /// owns the delete sweep + write). This variant lets
    /// [`MultiPreferenceVectors`](crate::entities::MultiPreferenceVectors)'s
    /// unified checkpoint fold cold-start users into the same atomic batch as the
    /// cluster rows, so a single swap covers both tiers. The `skip` predicate lets
    /// that caller exclude any user already written as a warm cluster row, so a
    /// legacy row can never collide with (and clobber) a warm row under the shared
    /// `Tag::Preference` key.
    pub fn append_legacy_rows(
        &self,
        batch: &mut crate::storage::WriteBatch,
        skip: impl Fn(u64) -> bool,
    ) {
        use crate::{
            schema::EntityId,
            storage::{Tag, encode_key},
        };
        for entry in &self.inner {
            let user_id = *entry.key();
            if skip(user_id) {
                continue;
            }
            let value = encode_legacy_row(self.update_count(user_id), entry.value());
            let key = encode_key(EntityId::new(0), Tag::Preference, &user_id.to_be_bytes());
            batch.put(key, value);
        }
    }

    /// Number of users with stored preferences.
    #[must_use]
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    /// Whether no preferences are stored.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    /// Get the current update count for a user (number of times `update` has been called).
    #[must_use]
    pub fn update_count(&self, user_id: u64) -> u64 {
        self.update_counts.get(&user_id).map_or(0, |c| *c)
    }

    /// Compute the adaptive learning rate for a user at their current update count.
    ///
    /// Formula: `base_alpha / (1 + ln(count + 1))`.
    #[must_use]
    #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
    pub fn adaptive_lr(&self, user_id: u64) -> f32 {
        let count = self.update_count(user_id);
        (f64::from(self.learning_rate) / (1.0 + (count as f64).ln_1p())) as f32
    }
}

// ── Checkpoint / Restore ─────────────────────────────────────────────────────

impl PreferenceVectors {
    /// Checkpoint every per-user preference vector + update count to durable
    /// storage under `Tag::Preference`, one row per user.
    ///
    /// All rows share the sentinel entity id (`0`) so [`restore`](Self::restore)
    /// can target them with a single prefix scan. The swap is staged into one
    /// `WriteBatch` and committed atomically (matching the co-engagement /
    /// community checkpoints): stale rows are deleted and current rows written in
    /// the same batch, so a crash mid-checkpoint never leaves a half-written set.
    ///
    /// Key suffix: `[user: 8B BE]`. Value: `[update_count: 8B LE][dim: 4B LE][f32 * dim LE]`.
    ///
    /// **In-engine this is superseded by
    /// [`MultiPreferenceVectors`](crate::entities::MultiPreferenceVectors), which
    /// owns all of `Tag::Preference` and writes cold-start users via
    /// `append_legacy_rows`.** This standalone path remains for direct single-vector
    /// embedders; it shares the `encode_legacy_row` writer so the two cannot drift.
    ///
    /// # Errors
    ///
    /// Returns storage errors from the underlying engine.
    pub fn checkpoint(&self, storage: &dyn crate::storage::StorageEngine) -> crate::Result<()> {
        use crate::{
            schema::EntityId,
            storage::{Tag, WriteBatch, entity_tag_prefix},
        };

        let prefix = entity_tag_prefix(EntityId::new(0), Tag::Preference);
        let mut batch = WriteBatch::with_capacity(self.inner.len() + 1);

        // Stage deletion of any stale rows (e.g. a user whose vector was cleared)
        // so the post-swap snapshot holds exactly the current set.
        for item in storage.scan_prefix(&prefix) {
            let (key, _) = item.map_err(crate::schema::TidalError::from)?;
            batch.delete(key);
        }

        // Write every current row through the single-sourced legacy encoder.
        self.append_legacy_rows(&mut batch, |_| false);

        storage
            .write_batch(batch)
            .map_err(crate::schema::TidalError::from)?;
        Ok(())
    }

    /// Restore per-user preference vectors + update counts from a
    /// `Tag::Preference` checkpoint. Skips rows whose stored dimension does not
    /// match this store's `dim` (a schema change) and rows that are torn.
    ///
    /// Each restored vector is re-normalized to unit length at the load boundary
    /// (after zeroing any NaN component) so the cosine-scoring invariant holds
    /// even for a torn or tampered row — see the inline comment at the
    /// normalization site.
    ///
    /// # Errors
    ///
    /// Returns storage errors from the underlying engine.
    pub fn restore(&self, storage: &dyn crate::storage::StorageEngine) -> crate::Result<()> {
        use crate::{
            schema::EntityId,
            storage::{Tag, entity_tag_prefix, parse_key},
        };

        let prefix = entity_tag_prefix(EntityId::new(0), Tag::Preference);
        let mut count = 0u64;
        for entry in storage.scan_prefix(&prefix) {
            let (key, value) = entry.map_err(crate::schema::TidalError::from)?;
            let Some((_, Tag::Preference, suffix)) = parse_key(&key) else {
                continue;
            };
            if suffix.len() < 8 {
                continue;
            }
            let user_id = u64::from_be_bytes(suffix[0..8].try_into().unwrap_or([0u8; 8]));
            // Decode + NaN-neutralize + re-normalize at the load boundary via the
            // single-sourced decoder (skips torn / dimension-mismatched rows). The
            // re-normalization re-establishes the unit-length invariant that
            // `cosine_similarity` assumes, so a torn row degrades to a zero
            // preference rather than a poisoned non-unit one.
            let Some((update_count, vec)) = decode_legacy_row(&value, self.dim) else {
                continue;
            };
            // Seed the adaptive update count so the learning rate resumes where
            // it was at checkpoint time.
            self.inner.insert(user_id, vec);
            self.update_counts.insert(user_id, update_count);
            let _ = EntityId::new(user_id);
            count += 1;
        }
        if count > 0 {
            tracing::info!(users = count, "preference vectors restored from checkpoint");
        }
        Ok(())
    }
}

/// L2-normalize an accumulated preference centroid in place, tolerating a cold
/// (all-zero) vector.
///
/// This is the **one** place the centroid zero-tolerance policy is named. The math
/// itself lives in [`crate::storage::vector::l2_normalize_in_place`]; this wrapper
/// only decides what a zero norm *means* here:
///
/// A preference centroid is legitimately all-zero until the first signal folds in,
/// so "no direction yet" is an ordinary state, not an error — the vector is left at
/// zero and will normalize on the next fold. Embedding writes deliberately do **not**
/// get this leniency (`storage::vector::lifecycle::ops` propagates the error) because
/// a zero embedding entering the HNSW index is unrecoverable garbage.
///
/// Both `preference` and `multi_preference` call this; before, each carried its own
/// copy of the arithmetic with a zero threshold ~2900x looser than the canonical one,
/// which is how the two directions of normalization drifted apart in the first place.
/// Vectors with `0 < ||v|| < 3.45e-4` are now left alone rather than amplified, since
/// dividing by a norm that small yields a direction made of rounding error.
pub(crate) fn normalize_centroid(vec: &mut [f32]) {
    if let Err(err) = l2_normalize_in_place(vec) {
        // A cold or numerically-zero centroid is the ONLY error this can return
        // (see `l2_norm_checked`), and it is an ordinary state here: leave the
        // vector exactly as it was.
        //
        // Anything else is unreachable today. Leaving the centroid untouched stays
        // the safe default if normalization ever grows a second failure mode, but
        // say so out loud rather than discarding it silently.
        if !matches!(err, VectorError::ZeroNormVector) {
            tracing::warn!(error = %err, "unexpected error normalizing preference centroid");
        }
    }
}

/// Canonical legacy single-vector row encoder: `[count:8 LE][dim:4 LE][f32*dim]`.
///
/// This is the **one** place the legacy row is written. `append_legacy_rows`, the
/// standalone [`PreferenceVectors::checkpoint`], and (via `decode_legacy_row`) the
/// multi-vector cold-start path all go through this pair so the on-disk format
/// cannot drift across writers/readers.
#[must_use]
pub(crate) fn encode_legacy_row(count: u64, vec: &[f32]) -> Vec<u8> {
    let mut value = Vec::with_capacity(8 + 4 + vec.len() * 4);
    value.extend_from_slice(&count.to_le_bytes());
    #[allow(clippy::cast_possible_truncation)]
    value.extend_from_slice(&(vec.len() as u32).to_le_bytes());
    for v in vec {
        value.extend_from_slice(&v.to_le_bytes());
    }
    value
}

/// Canonical legacy single-vector row decoder. Returns `(update_count,
/// normalized_vector)`, or `None` if the row is torn or its stored dimension does
/// not match `expected_dim`. Each component is NaN-neutralized and the vector is
/// re-normalized at the load boundary (the cosine-scoring invariant), mirroring
/// the restore contract. Paired with [`encode_legacy_row`]; both
/// [`PreferenceVectors::restore`] and the multi-vector store's legacy decode
/// delegate here so the format has exactly one encoder and one decoder.
#[must_use]
pub(crate) fn decode_legacy_row(value: &[u8], expected_dim: usize) -> Option<(u64, Vec<f32>)> {
    if value.len() < 12 {
        return None;
    }
    let update_count = u64::from_le_bytes(value[0..8].try_into().ok()?);
    let dim = u32::from_le_bytes(value[8..12].try_into().ok()?) as usize;
    if dim != expected_dim || value.len() < 12 + dim * 4 {
        return None;
    }
    let mut vec = Vec::with_capacity(dim);
    for i in 0..dim {
        let off = 12 + i * 4;
        let f = f32::from_le_bytes(value[off..off + 4].try_into().ok()?);
        vec.push(if f.is_nan() { 0.0 } else { f });
    }
    normalize_centroid(&mut vec);
    Some((update_count, vec))
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::float_cmp)]
mod tests {
    use super::*;

    #[test]
    fn set_and_get() {
        let pv = PreferenceVectors::new(3);
        assert!(pv.set(1, vec![3.0, 4.0, 0.0]));
        let v = pv.get(1).unwrap();
        // 3/5, 4/5, 0
        assert!((v[0] - 0.6).abs() < 1e-6);
        assert!((v[1] - 0.8).abs() < 1e-6);
        assert!((v[2] - 0.0).abs() < 1e-6);
    }

    #[test]
    fn set_wrong_dim_rejected() {
        let pv = PreferenceVectors::new(3);
        assert!(!pv.set(1, vec![1.0, 2.0]));
        assert!(pv.get(1).is_none());
    }

    #[test]
    fn update_creates_initial() {
        let pv = PreferenceVectors::new(3);
        assert!(pv.update(1, &[1.0, 0.0, 0.0]));
        let v = pv.get(1).unwrap();
        assert!((v[0] - 1.0).abs() < 1e-6);
    }

    #[test]
    fn update_blends() {
        let pv = PreferenceVectors::with_learning_rate(2, 0.5);
        let _ = pv.set(1, vec![1.0, 0.0]);
        let _ = pv.update(1, &[0.0, 1.0]);
        let v = pv.get(1).unwrap();
        // After blend: (0.5, 0.5), normalized: (1/sqrt(2), 1/sqrt(2))
        let expected = 1.0 / 2.0f32.sqrt();
        assert!((v[0] - expected).abs() < 1e-5);
        assert!((v[1] - expected).abs() < 1e-5);
    }

    #[test]
    fn cosine_similarity_normalized() {
        let pv = PreferenceVectors::new(3);
        let _ = pv.set(1, vec![1.0, 0.0, 0.0]);
        // Cosine with self = 1.0
        let sim = pv.cosine_similarity(1, &[1.0, 0.0, 0.0]).unwrap();
        assert!((sim - 1.0).abs() < 1e-6);
        // Orthogonal = 0.0
        let sim = pv.cosine_similarity(1, &[0.0, 1.0, 0.0]).unwrap();
        assert!(sim.abs() < 1e-6);
    }

    #[test]
    fn cosine_similarity_no_pref() {
        let pv = PreferenceVectors::new(3);
        assert!(pv.cosine_similarity(1, &[1.0, 0.0, 0.0]).is_none());
    }

    #[test]
    fn l2_normalize_zero_vec() {
        let mut v = vec![0.0f32, 0.0, 0.0];
        normalize_centroid(&mut v);
        assert!(v.iter().all(|&x| x == 0.0));
    }

    /// Pins the ONE behavior change from collapsing the three `l2_normalize`
    /// copies into `storage::vector::l2_normalize_in_place`.
    ///
    /// The deleted in-place copies rejected only `||v|| <= f32::EPSILON` (1.2e-7);
    /// the canonical implementation rejects `||v||^2 < f32::EPSILON`, i.e.
    /// `||v|| < 3.45e-4` — about 2900x stricter. A centroid in that gap is now left
    /// alone instead of being scaled to unit length.
    ///
    /// That is deliberate: dividing by a norm that small amplifies float noise by
    /// more than 2900x, so the "direction" produced is mostly rounding error.
    /// Leaving the vector as-is is the honest answer, and it matches what embeddings
    /// have always done. Asserted rather than merely documented so the choice cannot
    /// be un-made silently.
    #[test]
    fn normalize_centroid_leaves_numerically_zero_vector_untouched() {
        // ||v|| = 1e-5 * sqrt(3) = 1.73e-5, inside the gap between the two thresholds.
        let mut v = vec![1e-5f32, 1e-5, 1e-5];
        let before = v.clone();
        normalize_centroid(&mut v);
        assert_eq!(
            v, before,
            "a centroid with norm below the canonical threshold must be left as-is, \
             not amplified into a direction made of rounding error"
        );

        // Just above the threshold it must still normalize, so the guard is a floor
        // and not a silent no-op for real data.
        let mut real = vec![1.0f32, 2.0, 3.0];
        normalize_centroid(&mut real);
        let norm: f32 = real.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (1.0 - norm).abs() < 1e-5,
            "real centroid must reach unit norm, got {norm}"
        );
    }

    #[test]
    fn update_with_custom_rate_blends_at_given_lr() {
        let pv = PreferenceVectors::new(2);
        let _ = pv.set(1, vec![1.0, 0.0]);
        // lr=0.5 means 50% blend.
        let ok = pv.update_with_custom_rate(1, &[0.0, 1.0], 0.5);
        assert!(ok);
        let v = pv.get(1).unwrap();
        let expected = 1.0 / 2.0f32.sqrt();
        assert!((v[0] - expected).abs() < 1e-5);
        assert!((v[1] - expected).abs() < 1e-5);
    }

    #[test]
    fn update_with_custom_rate_wrong_dim() {
        let pv = PreferenceVectors::new(3);
        assert!(!pv.update_with_custom_rate(1, &[1.0], 0.1));
    }

    #[test]
    fn len_and_is_empty() {
        let pv = PreferenceVectors::new(3);
        assert!(pv.is_empty());
        assert_eq!(pv.len(), 0);
        let _ = pv.set(1, vec![1.0, 0.0, 0.0]);
        assert!(!pv.is_empty());
        assert_eq!(pv.len(), 1);
    }

    #[test]
    fn adaptive_lr_decays_monotonically() {
        let pv = PreferenceVectors::new(4);
        let embedding = [1.0f32, 0.0, 0.0, 0.0];

        // At count=0: alpha = 0.1 / (1 + ln(1)) = 0.1
        let alpha_at_0 = pv.adaptive_lr(1);
        assert!(
            (alpha_at_0 - 0.1).abs() < 1e-5,
            "alpha at count=0 should be 0.1, got {alpha_at_0}"
        );

        // Apply 1000 updates.
        for _ in 0..1000 {
            let _ = pv.update(1, &embedding);
        }

        let alpha_at_1000 = pv.adaptive_lr(1);

        // After 1000 updates, alpha should be substantially smaller.
        assert!(
            alpha_at_1000 < 0.05,
            "alpha at 1000 updates ({alpha_at_1000}) should be < 0.05"
        );
        assert!(
            alpha_at_1000 < alpha_at_0,
            "alpha must decay with update count"
        );

        // Verify intermediate counts also decay.
        let pv2 = PreferenceVectors::new(4);
        for _ in 0..100 {
            let _ = pv2.update(2, &embedding);
        }
        let alpha_at_100 = pv2.adaptive_lr(2);
        assert!(
            alpha_at_100 < alpha_at_0,
            "alpha at 100 should be < alpha at 0"
        );
        assert!(
            alpha_at_1000 < alpha_at_100,
            "alpha at 1000 should be < alpha at 100"
        );
    }

    /// Verify the decay formula using actual vector shift magnitudes.
    ///
    /// The AC states that after 1000 preference updates, new signals shift the
    /// vector less than at update count=1. This test measures the L2-norm of
    /// `pref_after - pref_before` at different update counts using an orthogonal
    /// interaction direction to maximize the visible shift.
    ///
    /// The `ln_1p` decay formula: `alpha = base / (1 + ln(count + 1))` achieves
    /// ~10–15% of the initial shift magnitude at count=1000 (the formula is
    /// logarithmic, not exponential; it cannot achieve < 5% in finite update counts
    /// without an unreasonably small `base_alpha`).
    #[test]
    fn adaptive_lr_shift_magnitude_decays_with_update_count() {
        // Use a fresh instance so we control update counts precisely.
        let pv = PreferenceVectors::new(4);

        // Helper to measure the L2 shift magnitude of one update in an
        // orthogonal direction from the current preference.
        let measure_shift = |pv: &PreferenceVectors, interaction: &[f32]| -> f32 {
            let before = pv.get(1).unwrap_or_else(|| vec![1.0, 0.0, 0.0, 0.0]);
            let _ = pv.update(1, interaction);
            let after = pv.get(1).unwrap();
            before
                .iter()
                .zip(after.iter())
                .map(|(a, b)| (b - a).powi(2))
                .sum::<f32>()
                .sqrt()
        };

        // Establish a baseline preference in direction [1, 0, 0, 0].
        let primary = [1.0f32, 0.0, 0.0, 0.0];
        let orthogonal = [0.0f32, 1.0, 0.0, 0.0];

        // Initial state via set (does not increment update_count).
        assert!(pv.set(1, primary.to_vec()));

        // Shift at count=0 (first adaptive update).
        let shift_at_0 = measure_shift(&pv, &orthogonal);
        // Reset back to primary to isolate shift measurement.
        assert!(pv.set(1, primary.to_vec()));

        // Apply 99 more updates to reach count=100.
        for _ in 0..99 {
            let _ = pv.update(1, &primary);
        }
        let shift_at_100 = measure_shift(&pv, &orthogonal);
        assert!(pv.set(1, primary.to_vec()));

        // Apply another 900 updates to reach count=1000.
        for _ in 0..900 {
            let _ = pv.update(1, &primary);
        }
        let shift_at_1000 = measure_shift(&pv, &orthogonal);

        // Verify monotonic decay of shift magnitude.
        assert!(
            shift_at_100 < shift_at_0,
            "shift at 100 ({shift_at_100:.6}) must be < shift at 0 ({shift_at_0:.6})"
        );
        assert!(
            shift_at_1000 < shift_at_100,
            "shift at 1000 ({shift_at_1000:.6}) must be < shift at 100 ({shift_at_100:.6})"
        );

        // Verify the shift at 1000 is substantially smaller than at count=0.
        // The ln_1p formula achieves ~10-15% of initial shift at count=1000.
        let ratio = shift_at_1000 / shift_at_0;
        assert!(
            ratio < 0.20,
            "shift at 1000 ({shift_at_1000:.6}) should be < 20% of shift at 0 ({shift_at_0:.6}); got {ratio:.3}"
        );
    }

    #[test]
    fn update_count_tracks_calls() {
        let pv = PreferenceVectors::new(2);
        assert_eq!(pv.update_count(1), 0);
        let _ = pv.update(1, &[1.0, 0.0]);
        assert_eq!(pv.update_count(1), 1);
        let _ = pv.update(1, &[0.0, 1.0]);
        assert_eq!(pv.update_count(1), 2);
        // Different user is independent.
        assert_eq!(pv.update_count(2), 0);
    }

    #[test]
    fn custom_rate_does_not_use_adaptive_lr() {
        let pv = PreferenceVectors::new(2);
        // Set a known state.
        let _ = pv.set(1, vec![1.0, 0.0]);
        // Use custom rate -- should NOT increment update_counts.
        let _ = pv.update_with_custom_rate(1, &[0.0, 1.0], 0.5);
        assert_eq!(
            pv.update_count(1),
            0,
            "update_with_custom_rate should not increment adaptive count"
        );
    }

    /// Encode a `Tag::Preference` checkpoint value: `[count:8 LE][dim:4 LE][f32*dim]`.
    fn encode_pref_value(count: u64, vec: &[f32]) -> Vec<u8> {
        let mut value = Vec::with_capacity(8 + 4 + vec.len() * 4);
        value.extend_from_slice(&count.to_le_bytes());
        value.extend_from_slice(&u32::try_from(vec.len()).unwrap().to_le_bytes());
        for v in vec {
            value.extend_from_slice(&v.to_le_bytes());
        }
        value
    }

    /// W12: a torn/tampered checkpoint row that is NOT unit-length must be
    /// re-normalized at the load boundary. `cosine_similarity` assumes the stored
    /// preference is unit-length (it divides only by the candidate norm), so a
    /// non-unit restored vector would mis-scale every similarity for that user.
    /// This covers both failure shapes: a NaN component (zeroed then non-unit) and
    /// a finite-but-non-unit row.
    #[test]
    fn restore_renormalizes_torn_rows_to_unit_length() {
        use crate::{
            schema::EntityId,
            storage::{InMemoryBackend, StorageEngine, Tag, encode_key},
        };

        let storage = InMemoryBackend::new();
        let dim = 3usize;

        // User 1: a NaN component. After zeroing, the residual [3,0,4] has norm 5,
        // i.e. NOT unit length — restore must normalize it to [0.6, 0, 0.8].
        let nan_row = encode_pref_value(7, &[3.0, f32::NAN, 4.0]);
        let key1 = encode_key(EntityId::new(0), Tag::Preference, &1u64.to_be_bytes());
        storage.put(&key1, &nan_row).unwrap();

        // User 2: a finite-but-non-unit row [0, 6, 8] (norm 10). Must normalize to
        // [0, 0.6, 0.8].
        let big_row = encode_pref_value(3, &[0.0, 6.0, 8.0]);
        let key2 = encode_key(EntityId::new(0), Tag::Preference, &2u64.to_be_bytes());
        storage.put(&key2, &big_row).unwrap();

        // User 3: an all-zero (fully-corrupt) row stays zero (degrades to a zero
        // preference, never a poisoned non-unit one).
        let zero_row = encode_pref_value(1, &[f32::NAN, f32::NAN, f32::NAN]);
        let key3 = encode_key(EntityId::new(0), Tag::Preference, &3u64.to_be_bytes());
        storage.put(&key3, &zero_row).unwrap();

        let pv = PreferenceVectors::new(dim);
        pv.restore(&storage).unwrap();

        let v1 = pv.get(1).expect("user 1 restored");
        let n1: f32 = v1.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (n1 - 1.0).abs() < 1e-5,
            "user 1 must be unit length, got {n1}"
        );
        assert!((v1[0] - 0.6).abs() < 1e-5 && (v1[2] - 0.8).abs() < 1e-5);
        assert_eq!(pv.update_count(1), 7, "update count must round-trip");

        let v2 = pv.get(2).expect("user 2 restored");
        let n2: f32 = v2.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (n2 - 1.0).abs() < 1e-5,
            "user 2 must be unit length, got {n2}"
        );

        let v3 = pv.get(3).expect("user 3 restored");
        assert!(
            v3.iter().all(|&x| x == 0.0),
            "a fully-corrupt row degrades to a zero preference, not a poisoned one"
        );

        // Cosine of the unit-normalized user-1 pref with its own direction is 1.0,
        // proving the invariant cosine_similarity depends on now holds.
        let sim = pv.cosine_similarity(1, &[0.6, 0.0, 0.8]).unwrap();
        assert!(
            (sim - 1.0).abs() < 1e-5,
            "cosine self-similarity must be 1.0, got {sim}"
        );
    }

    mod proptests {
        use proptest::prelude::*;

        use super::*;

        proptest! {
            /// After any sequence of updates, the L2 norm stays approximately 1.0.
            #[test]
            fn l2_norm_invariant(
                updates in proptest::collection::vec(
                    proptest::collection::vec(-1.0f32..1.0f32, 4..=4),
                    1..20
                ),
            ) {
                let pv = PreferenceVectors::new(4);
                for emb in &updates {
                    let _ = pv.update(1, emb);
                }
                let v = pv.get(1).unwrap();
                let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
                // After a sequence of updates, the vector should be unit-length
                // (within floating-point tolerance) or exactly zero if all inputs
                // collapse to the origin.
                prop_assert!(
                    (norm - 1.0).abs() < 1e-4 || norm < f32::EPSILON,
                    "norm was {norm}, expected ~1.0"
                );
            }
        }
    }
}