keydr/src/engine/key_stats.rs

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

const EMA_ALPHA: f64 = 0.1;
const DEFAULT_TARGET_CPM: f64 = 175.0;

/// Sticky-mastery promotion thresholds. A key is promoted to `mastered: true`
/// only when it meets all three gates on a ranked stats update.
pub const MASTERY_MIN_SPEED_CONFIDENCE: f64 = 1.05;
pub const MASTERY_MIN_SAMPLES: usize = 20;
pub const MASTERY_MAX_ERROR_RATE_EMA: f64 = 0.05;

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct KeyStat {
    pub filtered_time_ms: f64,
    pub best_time_ms: f64,
    pub confidence: f64,
    pub sample_count: usize,
    pub recent_times: Vec<f64>,
    #[serde(default)]
    pub error_count: usize,
    #[serde(default)]
    pub total_count: usize,
    #[serde(default = "default_error_rate_ema")]
    pub error_rate_ema: f64,
    // Sticky mastery: once true, never reset by normal drill updates.
    // Only authoritative when read from `ranked_key_stats` (unranked `key_stats`
    // stores share the type but are not allowed to drive progression).
    #[serde(default)]
    pub mastered: bool,
}

fn default_error_rate_ema() -> f64 {
    0.5
}

impl KeyStat {
    /// Returns true if this stat currently satisfies every sticky-mastery gate.
    /// Does not mutate anything — callers promote via `KeyStatsStore` helpers.
    pub fn qualifies_for_mastery(&self) -> bool {
        self.confidence >= MASTERY_MIN_SPEED_CONFIDENCE
            && self.sample_count >= MASTERY_MIN_SAMPLES
            && self.error_rate_ema <= MASTERY_MAX_ERROR_RATE_EMA
    }
}

impl Default for KeyStat {
    fn default() -> Self {
        Self {
            filtered_time_ms: 1000.0,
            best_time_ms: f64::MAX,
            confidence: 0.0,
            sample_count: 0,
            recent_times: Vec::new(),
            error_count: 0,
            total_count: 0,
            error_rate_ema: 0.5,
            mastered: false,
        }
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct KeyStatsStore {
    pub stats: HashMap<char, KeyStat>,
    pub target_cpm: f64,
}

impl Default for KeyStatsStore {
    fn default() -> Self {
        Self {
            stats: HashMap::new(),
            target_cpm: DEFAULT_TARGET_CPM,
        }
    }
}

impl KeyStatsStore {
    pub fn update_key(&mut self, key: char, time_ms: f64) {
        let stat = self.stats.entry(key).or_default();
        stat.sample_count += 1;
        stat.total_count += 1;

        if stat.sample_count == 1 {
            stat.filtered_time_ms = time_ms;
        } else {
            stat.filtered_time_ms = EMA_ALPHA * time_ms + (1.0 - EMA_ALPHA) * stat.filtered_time_ms;
        }

        stat.best_time_ms = stat.best_time_ms.min(stat.filtered_time_ms);

        let target_time_ms = 60000.0 / self.target_cpm;
        stat.confidence = target_time_ms / stat.filtered_time_ms;

        stat.recent_times.push(time_ms);
        if stat.recent_times.len() > 30 {
            stat.recent_times.remove(0);
        }

        // Update error rate EMA (correct stroke = 0.0 signal)
        if stat.total_count == 1 {
            stat.error_rate_ema = 0.0;
        } else {
            stat.error_rate_ema = EMA_ALPHA * 0.0 + (1.0 - EMA_ALPHA) * stat.error_rate_ema;
        }
    }

    pub fn get_confidence(&self, key: char) -> f64 {
        self.stats.get(&key).map(|s| s.confidence).unwrap_or(0.0)
    }

    #[allow(dead_code)]
    pub fn get_stat(&self, key: char) -> Option<&KeyStat> {
        self.stats.get(&key)
    }

    /// Record an error for a key (increments error_count and total_count).
    /// Does NOT update timing/confidence (those are only updated for correct strokes).
    pub fn update_key_error(&mut self, key: char) {
        let stat = self.stats.entry(key).or_default();
        stat.error_count += 1;
        stat.total_count += 1;

        // Update error rate EMA (error stroke = 1.0 signal)
        if stat.total_count == 1 {
            stat.error_rate_ema = 1.0;
        } else {
            stat.error_rate_ema = EMA_ALPHA * 1.0 + (1.0 - EMA_ALPHA) * stat.error_rate_ema;
        }
    }

    /// Ranked counterpart of [`update_key`]. Accumulates identical metrics and
    /// additionally evaluates sticky-mastery promotion.
    ///
    /// Promotion lives only on the ranked methods: `ranked_key_stats` is the sole
    /// authority for progression, so unranked/partial activity must never set the
    /// `mastered` bit even when its metrics would satisfy the gate.
    pub fn update_key_ranked(&mut self, key: char, time_ms: f64) {
        self.update_key(key, time_ms);
        self.promote_if_qualified(key);
    }

    /// Ranked counterpart of [`update_key_error`]. See [`update_key_ranked`] for
    /// why promotion is ranked-only.
    ///
    /// Errors never unset sticky mastery, but promotion is still evaluated here so
    /// both ranked entry points share one promotion site (a key sitting right on
    /// the gate can still promote from already-accumulated samples/confidence).
    pub fn update_key_error_ranked(&mut self, key: char) {
        self.update_key_error(key);
        self.promote_if_qualified(key);
    }

    /// Single source of truth for promotion. Flips `mastered` from false to true
    /// when all thresholds are met. Never unsets mastery.
    fn promote_if_qualified(&mut self, key: char) {
        if let Some(stat) = self.stats.get_mut(&key)
            && !stat.mastered
            && stat.qualifies_for_mastery()
        {
            stat.mastered = true;
        }
    }

    /// Every key currently holding the sticky mastery bit.
    pub fn mastered_keys(&self) -> Vec<char> {
        self.stats
            .iter()
            .filter(|(_, s)| s.mastered)
            .map(|(&k, _)| k)
            .collect()
    }

    /// Re-apply a previously earned sticky mastery bit, creating a bare entry if
    /// the key has no stats left. Used when rebuilding from a history window that
    /// no longer contains the strokes that originally earned mastery — mastery is
    /// one-way, so a rebuild must never demote a key.
    pub fn seed_mastered(&mut self, key: char) {
        self.stats.entry(key).or_default().mastered = true;
    }

    /// Returns the sticky persisted mastery bit for a key.
    pub fn is_mastered(&self, key: char) -> bool {
        self.stats.get(&key).map(|s| s.mastered).unwrap_or(false)
    }

    /// Returns true if the key's current stats would satisfy the promotion gate.
    /// Useful for tests/diagnostics; live updates promote automatically.
    #[allow(dead_code)]
    pub fn qualifies_for_mastery(&self, key: char) -> bool {
        self.stats
            .get(&key)
            .map(|s| s.qualifies_for_mastery())
            .unwrap_or(false)
    }

    /// Count of keys with sticky mastery set.
    #[allow(dead_code)]
    pub fn mastered_count(&self) -> usize {
        self.stats.values().filter(|s| s.mastered).count()
    }

    /// EMA-based error rate for a key.
    pub fn smoothed_error_rate(&self, key: char) -> f64 {
        match self.stats.get(&key) {
            Some(s) => s.error_rate_ema,
            None => 0.5,
        }
    }
}

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

    #[test]
    fn test_initial_confidence_is_zero() {
        let store = KeyStatsStore::default();
        assert_eq!(store.get_confidence('a'), 0.0);
    }

    #[test]
    fn test_update_key_creates_stat() {
        let mut store = KeyStatsStore::default();
        store.update_key('e', 300.0);
        assert!(store.get_confidence('e') > 0.0);
        assert_eq!(store.stats.get(&'e').unwrap().sample_count, 1);
    }

    #[test]
    fn test_ema_converges() {
        let mut store = KeyStatsStore::default();
        // Type key fast many times - confidence should increase
        for _ in 0..50 {
            store.update_key('t', 200.0);
        }
        let conf = store.get_confidence('t');
        // At 175 CPM target, target_time = 60000/175 = 342.8ms
        // With 200ms typing time, confidence = 342.8/200 = 1.71
        assert!(
            conf > 1.0,
            "confidence should be > 1.0 for fast typing, got {conf}"
        );
    }

    #[test]
    fn test_slow_typing_low_confidence() {
        let mut store = KeyStatsStore::default();
        for _ in 0..50 {
            store.update_key('a', 1000.0);
        }
        let conf = store.get_confidence('a');
        // target_time = 342.8ms, typing at 1000ms -> conf = 342.8/1000 = 0.34
        assert!(
            conf < 1.0,
            "confidence should be < 1.0 for slow typing, got {conf}"
        );
    }

    #[test]
    fn test_ema_error_rate_correct_strokes() {
        let mut store = KeyStatsStore::default();
        // All correct strokes → EMA should be 0.0 for first, stay near 0
        store.update_key('a', 200.0);
        assert!((store.smoothed_error_rate('a') - 0.0).abs() < f64::EPSILON);
        for _ in 0..10 {
            store.update_key('a', 200.0);
        }
        assert!(
            store.smoothed_error_rate('a') < 0.01,
            "All correct → EMA near 0"
        );
    }

    #[test]
    fn test_ema_error_rate_error_strokes() {
        let mut store = KeyStatsStore::default();
        // First stroke is error
        store.update_key_error('b');
        assert!((store.smoothed_error_rate('b') - 1.0).abs() < f64::EPSILON);
        // Follow with correct strokes → EMA decays
        for _ in 0..20 {
            store.update_key('b', 200.0);
        }
        let rate = store.smoothed_error_rate('b');
        assert!(
            rate < 0.15,
            "After 20 correct, EMA should be < 0.15, got {rate}"
        );
    }

    #[test]
    fn test_ema_error_rate_default_for_missing_key() {
        let store = KeyStatsStore::default();
        assert!((store.smoothed_error_rate('z') - 0.5).abs() < f64::EPSILON);
    }

    #[test]
    fn test_ema_error_rate_serde_default() {
        // Verify backward compat: deserializing old data without error_rate_ema gets 0.5
        let json = r#"{"filtered_time_ms":200.0,"best_time_ms":200.0,"confidence":1.0,"sample_count":10,"recent_times":[],"error_count":2,"total_count":10}"#;
        let stat: KeyStat = serde_json::from_str(json).unwrap();
        assert!((stat.error_rate_ema - 0.5).abs() < f64::EPSILON);
    }

    #[test]
    fn test_mastered_defaults_false_on_deserialize() {
        let json = r#"{"filtered_time_ms":200.0,"best_time_ms":200.0,"confidence":1.2,"sample_count":50,"recent_times":[],"error_count":0,"total_count":50,"error_rate_ema":0.0}"#;
        let stat: KeyStat = serde_json::from_str(json).unwrap();
        assert!(
            !stat.mastered,
            "mastered should default to false when absent from JSON"
        );
    }

    #[test]
    fn test_no_mastery_from_speed_alone_with_high_error_rate() {
        let mut store = KeyStatsStore::default();
        // Inject many errors first to drive error_rate_ema high.
        for _ in 0..30 {
            store.update_key_error('a');
        }
        // Then just enough fast correct strokes to satisfy speed + samples,
        // but not so many that the error-rate EMA decays below the threshold.
        for _ in 0..20 {
            store.update_key_ranked('a', 200.0);
        }
        let stat = store.get_stat('a').unwrap();
        assert!(
            stat.confidence >= MASTERY_MIN_SPEED_CONFIDENCE,
            "precondition: fast enough"
        );
        assert!(
            stat.sample_count >= MASTERY_MIN_SAMPLES,
            "precondition: enough samples"
        );
        // error_rate_ema decays with correct strokes but should still be above 0.05 here.
        assert!(
            stat.error_rate_ema > MASTERY_MAX_ERROR_RATE_EMA,
            "error_rate_ema should still exceed 0.05, got {}",
            stat.error_rate_ema
        );
        assert!(
            !store.is_mastered('a'),
            "key must not be mastered while error rate too high"
        );
    }

    #[test]
    fn test_no_mastery_with_too_few_samples() {
        let mut store = KeyStatsStore::default();
        // Few fast strokes (below MASTERY_MIN_SAMPLES)
        for _ in 0..(MASTERY_MIN_SAMPLES - 1) {
            store.update_key_ranked('b', 150.0);
        }
        let stat = store.get_stat('b').unwrap();
        assert!(stat.confidence >= MASTERY_MIN_SPEED_CONFIDENCE);
        assert!(stat.error_rate_ema <= MASTERY_MAX_ERROR_RATE_EMA);
        assert!(stat.sample_count < MASTERY_MIN_SAMPLES);
        assert!(
            !store.is_mastered('b'),
            "key must not be mastered without enough samples"
        );
    }

    #[test]
    fn test_mastery_when_all_thresholds_met() {
        let mut store = KeyStatsStore::default();
        for _ in 0..30 {
            store.update_key_ranked('c', 150.0);
        }
        let stat = store.get_stat('c').unwrap();
        assert!(stat.confidence >= MASTERY_MIN_SPEED_CONFIDENCE);
        assert!(stat.sample_count >= MASTERY_MIN_SAMPLES);
        assert!(stat.error_rate_ema <= MASTERY_MAX_ERROR_RATE_EMA);
        assert!(store.is_mastered('c'));
    }

    #[test]
    fn test_mastery_sticks_after_slow_strokes() {
        let mut store = KeyStatsStore::default();
        for _ in 0..30 {
            store.update_key_ranked('d', 150.0);
        }
        assert!(store.is_mastered('d'));
        // Many slow strokes — confidence will drop below 1.0 eventually.
        for _ in 0..50 {
            store.update_key_ranked('d', 2000.0);
        }
        let conf = store.get_confidence('d');
        assert!(conf < 1.0, "confidence should have dropped, got {conf}");
        assert!(
            store.is_mastered('d'),
            "sticky mastery must persist after slow strokes"
        );
    }

    #[test]
    fn test_mastery_sticks_after_errors() {
        let mut store = KeyStatsStore::default();
        for _ in 0..30 {
            store.update_key_ranked('e', 150.0);
        }
        assert!(store.is_mastered('e'));
        for _ in 0..10 {
            store.update_key_error_ranked('e');
        }
        assert!(
            store.is_mastered('e'),
            "sticky mastery must persist after later errors"
        );
    }

    #[test]
    fn test_mastered_count_tracks_promotions() {
        let mut store = KeyStatsStore::default();
        assert_eq!(store.mastered_count(), 0);
        for _ in 0..30 {
            store.update_key_ranked('x', 150.0);
        }
        assert_eq!(store.mastered_count(), 1);
        for _ in 0..30 {
            store.update_key_ranked('y', 150.0);
        }
        assert_eq!(store.mastered_count(), 2);
    }

    #[test]
    fn test_unranked_updates_never_promote_mastery() {
        let mut unranked = KeyStatsStore::default();
        let mut ranked = KeyStatsStore::default();
        // Identical, comfortably qualifying activity applied to both stores.
        for _ in 0..30 {
            unranked.update_key('m', 150.0);
            ranked.update_key_ranked('m', 150.0);
        }

        assert!(
            unranked.qualifies_for_mastery('m'),
            "precondition: unranked metrics do satisfy the gate"
        );
        assert!(
            !unranked.is_mastered('m'),
            "unranked activity must never set the sticky bit — only ranked stats are authoritative"
        );
        assert!(
            ranked.is_mastered('m'),
            "the same activity on ranked stats must promote"
        );
        assert_eq!(unranked.mastered_count(), 0);
        assert_eq!(ranked.mastered_count(), 1);
    }

    #[test]
    fn test_unranked_error_updates_never_promote_mastery() {
        let mut unranked = KeyStatsStore::default();
        for _ in 0..30 {
            unranked.update_key('n', 150.0);
        }
        // An error update is the other promotion entry point; the unranked
        // variant must not promote on it either.
        unranked.update_key_error('n');
        assert!(
            !unranked.is_mastered('n'),
            "unranked error updates must not promote mastery"
        );
    }

    #[test]
    fn test_seed_mastered_restores_bit_without_stats() {
        let mut store = KeyStatsStore::default();
        store.seed_mastered('z');
        assert!(store.is_mastered('z'));
        assert_eq!(store.get_stat('z').unwrap().sample_count, 0);
        // Seeding must not be undone by subsequent ordinary activity.
        for _ in 0..10 {
            store.update_key_ranked('z', 2000.0);
        }
        assert!(store.is_mastered('z'));
    }

    #[test]
    fn test_mastered_keys_lists_only_promoted_keys() {
        let mut store = KeyStatsStore::default();
        for _ in 0..30 {
            store.update_key_ranked('p', 150.0);
        }
        for _ in 0..10 {
            store.update_key_ranked('q', 150.0);
        }
        let mut mastered = store.mastered_keys();
        mastered.sort_unstable();
        assert_eq!(mastered, vec!['p']);
    }
}