feedsignal/crates/core/src/affinity.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Per-topic affinity scores, in `[-1.0, 1.0]`, updated from observed
/// engagement vs. predicted relevance. Persisted as a single row per user
/// (JSON blob) in `feedsignal-db`; the event log remains the source of
/// truth and this can always be rebuilt by replaying it.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TopicAffinities {
    scores: HashMap<String, f64>,
}

/// How much a single feedback event moves an affinity score. Kept small so
/// no single article dominates a topic's long-run trend.
const LEARNING_RATE: f64 = 0.15;

/// Fraction of every affinity pulled back toward zero on each nightly decay
/// pass, so stale interests fade instead of anchoring the model forever.
const DAILY_DECAY: f64 = 0.02;

impl TopicAffinities {
    /// Get score for given topic
    pub fn score(&self, topic: &str) -> f64 {
        self.scores.get(topic).copied().unwrap_or(0.0)
    }

    /// Mean affinity across an article's topics; 0.0 for an untagged
    /// article (neutral, defers entirely to the embedding/LLM stages).
    pub fn get_mean_affinity(&self, topics: &[String]) -> f64 {
        if topics.is_empty() {
            return 0.0;
        }
        topics.iter().map(|t| self.score(t)).sum::<f64>() / topics.len() as f64
    }

    /// Update affinities for an article's topics from an observed
    /// `surprise`: `engagement_score - predicted_relevance_score`, both in
    /// `[0.0, 1.0]` (see `scoring::engagement_score`). Positive surprise
    /// (the user engaged more than the pipeline predicted) nudges those
    /// topics up; negative surprise nudges them down. This is what lets
    /// "the model thought this was irrelevant, but I read the whole thing"
    /// actually change future behavior.
    pub fn apply_feedback(&mut self, topics: &[String], surprise: f64) {
        for topic in topics {
            let current = self.score(topic);
            let updated = (current + LEARNING_RATE * surprise).clamp(-1.0, 1.0);
            self.scores.insert(topic.clone(), updated);
        }
    }

    /// Run once per day (see the scheduler in `feedsignal-web`) to let
    /// affinities the user hasn't reinforced recently drift back toward
    /// neutral rather than staying permanently pinned from a few old
    /// signals.
    pub fn decay(&mut self) {
        for v in self.scores.values_mut() {
            if v.abs() <= DAILY_DECAY {
                *v = 0.0;
            } else {
                *v -= v.signum() * DAILY_DECAY;
            }
        }
        self.scores.retain(|_, v| v.abs() > 1e-4);
    }

    pub fn top_n(&self, n: usize) -> Vec<(&str, f64)> {
        let mut items: Vec<_> = self.scores.iter().map(|(k, v)| (k.as_str(), *v)).collect();
        items.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        items.truncate(n);
        items
    }
}

/// Converts a raw reading interaction into an engagement score in
/// `[0.0, 1.0]`, and folds in the explicit star/dismiss signal.
///
/// - Never opened: 0.0
/// - Opened: `dwell_seconds / estimated_read_seconds`, capped at 1.0, so
///   skimming half an article scores lower than reading it fully.
/// - Starred: +0.3 on top (capped at 1.0) — an explicit "yes" beyond dwell
///   time alone.
/// - Dismissed without opening: -0.3, floored at 0.0's negative counterpart
///   handled by the caller via `surprise` (engagement itself never goes
///   negative; the *deviation* from a predicted score can).
pub fn engagement_score(
    opened: bool,
    dwell_seconds: Option<u32>,
    estimated_read_seconds: Option<u32>,
    starred: bool,
    dismissed: bool,
) -> f64 {
    if dismissed && !opened {
        return 0.0;
    }

    let mut score = if !opened {
        0.0
    } else {
        match (dwell_seconds, estimated_read_seconds) {
            (Some(dwell), Some(est)) if est > 0 => (dwell as f64 / est as f64).min(1.0),
            // Opened but we don't yet know dwell time / read-time estimate:
            // credit partial engagement rather than 0 or 1.
            _ => 0.5,
        }
    };

    if starred {
        score = (score + 0.3).min(1.0);
    }
    score
}

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

    // --- apply_feedback: new = clamp(current + LEARNING_RATE(0.15) * surprise, -1, 1) ---

    /// Positive surprise (engaged more than predicted) should move the
    /// score up, never down or unchanged.
    #[test]
    fn under_predicted_relevance_boosts_topic() {
        let mut aff = TopicAffinities::default();
        let topics = vec!["rust".to_string()];
        // Model predicted 0.2 relevance, user fully read it: surprise = 0.8.
        aff.apply_feedback(&topics, 0.8);
        assert!(aff.score("rust") > 0.0);
    }

    /// Negative surprise (engaged less than predicted) should move the
    /// score down, the mirror image of the boost case above.
    #[test]
    fn over_predicted_relevance_lowers_topic() {
        let mut aff = TopicAffinities::default();
        let topics = vec!["crypto".to_string()];
        // Model predicted 0.9, user dismissed unread: engagement 0, surprise = -0.9.
        aff.apply_feedback(&topics, -0.9);
        assert!(aff.score("crypto") < 0.0);
    }

    /// Scores are documented to live in [-1.0, 1.0]. Repeated max-surprise
    /// feedback would overshoot 1.0 without the clamp, so this guards the
    /// invariant directly rather than trusting a single update.
    #[test]
    fn apply_feedback_clamps_at_positive_one() {
        let mut aff = TopicAffinities::default();
        let topics = vec!["rust".to_string()];
        for _ in 0..20 {
            aff.apply_feedback(&topics, 1.0);
        }
        assert_eq!(aff.score("rust"), 1.0);
    }

    /// Verifies a decay from a large negative value doesn't overshoot and go beyond -1.0
    #[test]
    fn apply_feedback_clamps_at_negative_one() {
        let mut aff = TopicAffinities::default();
        let topics = vec!["crypto".to_string()];
        for _ in 0..20 {
            aff.apply_feedback(&topics, -1.0);
        }
        assert_eq!(aff.score("crypto"), -1.0);
    }

    /// apply_feedback loops over every topic on the article and applies
    /// the same surprise to each independently; it must not skip topics
    /// or bleed the update into topics the article wasn't tagged with.
    #[test]
    fn apply_feedback_updates_every_topic_on_the_article() {
        let mut aff = TopicAffinities::default();
        let topics = vec!["rust".to_string(), "async".to_string()];
        aff.apply_feedback(&topics, 0.4);
        assert_eq!(aff.score("rust"), 0.15 * 0.4);
        assert_eq!(aff.score("async"), 0.15 * 0.4);
        // Untouched topics are unaffected.
        assert_eq!(aff.score("crypto"), 0.0);
    }

    /// surprise = 0.0 means engagement exactly matched the prediction, so
    /// the score shouldn't move at all (current + 0.15 * 0.0 == current).
    #[test]
    fn apply_feedback_zero_surprise_is_a_noop() {
        let mut aff = TopicAffinities::default();
        let topics = vec!["rust".to_string()];
        aff.apply_feedback(&topics, 0.5);
        let before = aff.score("rust");
        aff.apply_feedback(&topics, 0.0);
        assert_eq!(aff.score("rust"), before);
    }

    // --- decay: v -= sign(v) * DAILY_DECAY(0.02), settling at 0 instead of overshooting ---

    /// Core decay behavior: a positive score should shrink toward zero
    /// after one nightly pass, without crossing it.
    #[test]
    fn decay_pulls_toward_zero() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["rust".to_string()], 1.0);
        let before = aff.score("rust");
        aff.decay();
        assert!(aff.score("rust") < before);
        assert!(aff.score("rust") > 0.0);
    }

    /// Pins the exact arithmetic (not just the direction) so a future
    /// change to the decay formula is caught immediately.
    /// surprise 1.0 -> 0.15, then one decay pass subtracts DAILY_DECAY (0.02).
    #[test]
    fn decay_gives_expected_value() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["java".to_string()], 1.0);
        aff.decay();

        assert_eq!(aff.score("java"), 0.13);
    }

    /// Regression test for a real bug: subtracting a fixed 0.02 from a
    /// smaller score (e.g. 0.015) used to flip its sign to -0.005 instead
    /// of landing on 0.0, which would make the score oscillate around
    /// zero on every subsequent decay pass rather than settling.
    #[test]
    fn decay_settles_at_zero_instead_of_overshooting() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["rust".to_string()], 0.1); // score = 0.015
        aff.decay();
        assert_eq!(aff.score("rust"), 0.0);
    }

    /// Same fix as above, verified on the negative side, and also checks
    /// that the post-decay prune (dropping |v| <= 1e-4) actually removes
    /// the entry rather than leaving a stray 0.0 in the map.
    #[test]
    fn decay_prunes_negative_scores_that_settle_at_zero() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["crypto".to_string()], -0.1); // score = -0.015
        aff.decay();
        assert_eq!(aff.score("crypto"), 0.0);
    }

    /// decay_pulls_toward_zero's mirror image: negative scores should
    /// shrink in magnitude too, not just positive ones.
    #[test]
    fn decay_is_symmetric_for_negative_scores() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["crypto".to_string()], -1.0);
        let before = aff.score("crypto");
        aff.decay();
        assert!(aff.score("crypto") > before);
        assert!(aff.score("crypto") < 0.0);
    }

    // --- score / get_mean_affinity ---

    /// A topic with no feedback yet must read as neutral (0.0), not
    /// panic or return some other sentinel.
    #[test]
    fn score_defaults_to_zero_for_unknown_topic() {
        let aff = TopicAffinities::default();
        assert_eq!(aff.score("never-seen"), 0.0);
    }

    /// Documented behavior for untagged articles: defer entirely to the
    /// embedding/LLM stages by returning a neutral 0.0 rather than
    /// dividing by zero.
    #[test]
    fn get_mean_affinity_given_empty_topics_returns_zero() {
        let aff = TopicAffinities::default();
        assert_eq!(aff.get_mean_affinity(&[]), 0.0);
    }

    /// Confirms it's a plain arithmetic mean: an equally strong positive
    /// and negative topic on the same article should cancel out to 0.0.
    #[test]
    fn get_mean_affinity_averages_across_topics() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["rust".to_string()], 1.0); // 0.15
        aff.apply_feedback(&["crypto".to_string()], -1.0); // -0.15
        let topics = vec!["rust".to_string(), "crypto".to_string()];
        assert_eq!(aff.get_mean_affinity(&topics), 0.0);
    }

    /// A topic mix of "known" and "never seen" shouldn't shrink the
    /// denominator or get skipped — the unscored topic counts as 0.0 in
    /// the average, per score()'s default.
    #[test]
    fn get_mean_affinity_treats_unscored_topics_as_zero() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["rust".to_string()], 1.0); // 0.15
        let topics = vec!["rust".to_string(), "never-seen".to_string()];
        assert_eq!(aff.get_mean_affinity(&topics), 0.075);
    }

    // --- top_n ---

    /// top_n is used to surface a user's strongest interests, so it must
    /// sort highest-first (not insertion order) and respect the limit.
    #[test]
    fn top_n_sorts_descending_and_truncates() {
        let mut aff = TopicAffinities::default();
        aff.apply_feedback(&["low".to_string()], 0.2);
        aff.apply_feedback(&["high".to_string()], 1.0);
        aff.apply_feedback(&["mid".to_string()], 0.5);

        let top = aff.top_n(2);
        assert_eq!(top.len(), 2);
        assert_eq!(top[0].0, "high");
        assert_eq!(top[1].0, "mid");
    }

    // --- engagement_score ---

    /// No signal at all (not opened, not dismissed) is neutral, not
    /// penalized.
    #[test]
    fn engagement_score_never_opened_is_zero() {
        assert_eq!(engagement_score(false, None, None, false, false), 0.0);
    }

    /// Dismissing without opening is an explicit negative signal, but
    /// engagement_score itself is floored at 0.0 (the doc comment notes
    /// the negative direction is expressed later via `surprise`, not
    /// here) — this pins that the dismissed+!opened branch returns 0.0,
    /// not a negative number.
    #[test]
    fn engagement_score_dismissed_without_opening_is_zero() {
        assert_eq!(engagement_score(false, None, None, false, true), 0.0);
    }

    /// Reading half the estimated time should score as half-engaged.
    #[test]
    fn engagement_score_opened_uses_dwell_over_estimate_ratio() {
        assert_eq!(
            engagement_score(true, Some(30), Some(60), false, false),
            0.5
        );
    }

    /// Dwelling far longer than the estimate (e.g. left the tab open)
    /// must not push the score above the documented [0.0, 1.0] range.
    #[test]
    fn engagement_score_opened_caps_ratio_at_one() {
        assert_eq!(
            engagement_score(true, Some(600), Some(60), false, false),
            1.0
        );
    }

    /// When we simply don't have dwell/estimate data yet, the code
    /// credits partial engagement (0.5) rather than assuming 0 (unfairly
    /// penalizing) or 1 (unfairly rewarding).
    #[test]
    fn engagement_score_opened_without_dwell_or_estimate_defaults_to_half() {
        assert_eq!(engagement_score(true, None, None, false, false), 0.5);
    }

    /// est == 0 would divide by zero, so the `est > 0` guard routes this
    /// case to the same "unknown read time" default (0.5) instead of
    /// panicking or producing NaN/infinity.
    #[test]
    fn engagement_score_opened_with_zero_estimate_defaults_to_half() {
        assert_eq!(engagement_score(true, Some(10), Some(0), false, false), 0.5);
    }

    /// Starring is an explicit "yes" beyond dwell time: it should add
    /// 0.3 on top of the dwell-ratio score.
    #[test]
    fn engagement_score_starred_adds_bonus() {
        assert_eq!(engagement_score(true, Some(30), Some(60), true, false), 0.8);
    }

    /// The +0.3 star bonus must also respect the 1.0 ceiling, even when
    /// the dwell ratio alone is already at the max.
    #[test]
    fn engagement_score_starred_bonus_caps_at_one() {
        assert_eq!(engagement_score(true, Some(60), Some(60), true, false), 1.0);
    }
}