//! 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.

use dashmap::DashMap;

/// 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;
        }
        l2_normalize(&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 {
        use dashmap::mapref::entry::Entry;

        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)]
        #[allow(clippy::cast_possible_truncation)]
        let lr = (f64::from(self.learning_rate) / (1.0 + (count as f64).ln_1p())) as f32;

        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);
                }
                l2_normalize(pref);
            }
            Entry::Vacant(vac) => {
                let mut v = interaction_embedding.to_vec();
                l2_normalize(&mut v);
                vac.insert(v);
            }
        }
        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 {
        use dashmap::mapref::entry::Entry;

        if interaction_embedding.len() != self.dim {
            return false;
        }

        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);
                }
                l2_normalize(pref);
            }
            Entry::Vacant(vac) => {
                let mut v = interaction_embedding.to_vec();
                l2_normalize(&mut v);
                vac.insert(v);
            }
        }
        true
    }

    /// 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)
    }

    /// 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
    }
}

/// L2-normalize a vector in-place. If the vector has zero magnitude, it remains
/// as-is (all zeros).
fn l2_normalize(vec: &mut [f32]) {
    let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > f32::EPSILON {
        for v in vec.iter_mut() {
            *v /= norm;
        }
    }
}

#[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];
        l2_normalize(&mut v);
        assert!(v.iter().all(|&x| x == 0.0));
    }

    #[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"
        );
    }

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

        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"
                );
            }
        }
    }
}
