tawasuyu/arje

1//! Observador estadístico. Mantiene marginales y co-ocurrencias dentro de una
2//! ventana deslizante. Calcula entropía de Shannon e información mutua para
3//! identificar correlaciones significativas.
4//!
5//! Diseño:
6//!   - Counters incrementales: cada `record()` es O(window_size) en el peor
7//!     caso (actualiza co-ocurrencias con cada evento del window).
8//!   - Sin recomputaciones globales: marginales y joint counts son state.
9//!   - El cálculo de H(X), P(B|A), I(A;B) es O(|distinct events|).
10
11use crate::rules::EventKind;
12use std::collections::{HashMap, VecDeque};
13use std::time::Instant;
14
15/// Evento timestamped. El timestamp se conserva para futuras políticas de
16/// expiración por tiempo (no sólo por count).
17#[derive(Debug, Clone)]
18pub struct TimedEvent {
19    pub kind: EventKind,
20    pub at: Instant,
21}
22
23pub struct Observer {
24    window: VecDeque<TimedEvent>,
25    window_size: usize,
26    marginal: HashMap<EventKind, u64>,
27    cooccur: HashMap<(EventKind, EventKind), u64>,
28    total: u64,
29    /// Last-seen timestamps para aplicar decay en query time. None = sin
30    /// time-decay (modo tradicional).
31    last_seen_marginal: HashMap<EventKind, Instant>,
32    last_seen_cooccur: HashMap<(EventKind, EventKind), Instant>,
33    /// Half-life del decay exponencial en segundos. None = sin decay
34    /// (las consultas devuelven los counts crudos).
35    half_life_secs: Option<f64>,
36}
37
38impl Observer {
39    pub fn new(window_size: usize) -> Self {
40        Self {
41            window: VecDeque::with_capacity(window_size),
42            window_size,
43            marginal: HashMap::new(),
44            cooccur: HashMap::new(),
45            total: 0,
46            last_seen_marginal: HashMap::new(),
47            last_seen_cooccur: HashMap::new(),
48            half_life_secs: None,
49        }
50    }
51
52    /// Activa decay exponencial con half-life en segundos. λ = ln(2)/half_life.
53    /// Aplicado en query time sobre los counts crudos usando last_seen.
54    pub fn with_half_life(mut self, half_life_secs: f64) -> Self {
55        if half_life_secs > 0.0 {
56            self.half_life_secs = Some(half_life_secs);
57        }
58        self
59    }
60
61    pub fn half_life(&self) -> Option<f64> { self.half_life_secs }
62
63    /// Registra un evento. Actualiza marginales y co-ocurrencias contra todo
64    /// evento aún en la ventana.
65    pub fn record(&mut self, kind: EventKind) {
66        let now = Instant::now();
67        let timed = TimedEvent { kind: kind.clone(), at: now };
68
69        // Co-ocurrencias: este evento con cada uno previo en ventana.
70        for w in &self.window {
71            let key = (w.kind.clone(), kind.clone());
72            *self.cooccur.entry(key.clone()).or_insert(0) += 1;
73            self.last_seen_cooccur.insert(key, now);
74        }
75
76        self.window.push_back(timed);
77        if self.window.len() > self.window_size {
78            self.window.pop_front();
79        }
80
81        *self.marginal.entry(kind.clone()).or_insert(0) += 1;
82        self.last_seen_marginal.insert(kind, now);
83        self.total += 1;
84    }
85
86    /// Aplica el decay sobre un count crudo dado el `last_seen` correspondiente.
87    /// Si half_life es None, devuelve el count tal cual (sin decay).
88    fn decay(&self, count: u64, last_seen: Option<Instant>) -> f64 {
89        let raw = count as f64;
90        let (hl, last) = match (self.half_life_secs, last_seen) {
91            (Some(hl), Some(t)) => (hl, t),
92            _ => return raw,
93        };
94        let age_secs = Instant::now().duration_since(last).as_secs_f64();
95        raw * 0.5_f64.powf(age_secs / hl)
96    }
97
98    /// Marginal con decay aplicado.
99    pub fn marginal_decayed(&self, k: &EventKind) -> f64 {
100        let raw = self.marginal.get(k).copied().unwrap_or(0);
101        let last = self.last_seen_marginal.get(k).copied();
102        self.decay(raw, last)
103    }
104
105    /// Cooccurrence con decay aplicado.
106    pub fn cooccur_decayed(&self, a: &EventKind, b: &EventKind) -> f64 {
107        let raw = self.cooccur.get(&(a.clone(), b.clone())).copied().unwrap_or(0);
108        let last = self.last_seen_cooccur.get(&(a.clone(), b.clone())).copied();
109        self.decay(raw, last)
110    }
111
112    /// Entropía de Shannon de la distribución marginal de eventos.
113    /// H(X) = −Σ p(x) log₂ p(x). Unidad: bits.
114    pub fn shannon_entropy(&self) -> f64 {
115        if self.total == 0 { return 0.0; }
116        let total = self.total as f64;
117        self.marginal.values()
118            .map(|&c| {
119                let p = c as f64 / total;
120                if p > 0.0 { -p * p.log2() } else { 0.0 }
121            })
122            .sum()
123    }
124
125    /// P(b | a) = "dado que algo siguió a `a` dentro del window, qué fracción
126    /// fue `b`". Suma 1 sobre todos los b posibles para un a fijo.
127    ///
128    /// Implementación: cooccur_decayed(a, b) / Σ_x cooccur_decayed(a, x).
129    /// Si half_life is None, los decayed values son los counts crudos.
130    pub fn conditional_prob(&self, a: &EventKind, b: &EventKind) -> f64 {
131        let joint = self.cooccur_decayed(a, b);
132        let row_total: f64 = self.cooccur.keys()
133            .filter(|(x, _)| x == a)
134            .map(|(x, y)| self.cooccur_decayed(x, y))
135            .sum();
136        if row_total <= 0.0 { 0.0 } else { joint / row_total }
137    }
138
139    /// Información mutua puntual entre `a` y `b` con decay aplicado:
140    /// PMI(a, b) = log₂( P(a, b) / (P(a) · P(b)) ).
141    /// Positivo → más correlacionados de lo que sugiere independencia.
142    pub fn pmi(&self, a: &EventKind, b: &EventKind) -> f64 {
143        // Total decayed: suma de marginales con decay (no usamos self.total
144        // directo porque debería ser consistente con los decayed values).
145        let total_decayed: f64 = self.marginal.keys()
146            .map(|k| self.marginal_decayed(k))
147            .sum();
148        if total_decayed <= 0.0 { return 0.0; }
149        let joint = self.cooccur_decayed(a, b) / total_decayed;
150        let pa = self.marginal_decayed(a) / total_decayed;
151        let pb = self.marginal_decayed(b) / total_decayed;
152        if joint <= 0.0 || pa <= 0.0 || pb <= 0.0 { return 0.0; }
153        (joint / (pa * pb)).log2()
154    }
155
156    /// Información mutua acumulada de la pareja (a, b) ponderada por su
157    /// probabilidad conjunta. Útil como medida de "interés" del par.
158    pub fn weighted_pmi(&self, a: &EventKind, b: &EventKind) -> f64 {
159        if self.total == 0 { return 0.0; }
160        let joint = self.cooccur
161            .get(&(a.clone(), b.clone()))
162            .copied()
163            .unwrap_or(0) as f64 / self.total as f64;
164        joint * self.pmi(a, b)
165    }
166
167    pub fn marginals(&self) -> &HashMap<EventKind, u64> { &self.marginal }
168    pub fn cooccurrences(&self) -> &HashMap<(EventKind, EventKind), u64> { &self.cooccur }
169    pub fn total(&self) -> u64 { self.total }
170    pub fn window_size(&self) -> usize { self.window_size }
171    pub fn current_window(&self) -> usize { self.window.len() }
172
173    /// Últimos N eventos del window, en orden cronológico (más viejo primero).
174    /// Si N > window.len(), devuelve todo el window.
175    pub fn recent(&self, n: usize) -> impl Iterator<Item = &TimedEvent> {
176        let start = self.window.len().saturating_sub(n);
177        self.window.range(start..)
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184    use crate::rules::EventKind::*;
185
186    #[test]
187    fn entropy_zero_for_single_event() {
188        let mut obs = Observer::new(10);
189        for _ in 0..5 { obs.record(EnteSpawned); }
190        // Distribución degenerada: una sola observación posible → H = 0.
191        assert!(obs.shannon_entropy() < 1e-9);
192    }
193
194    #[test]
195    fn entropy_one_for_balanced_binary() {
196        let mut obs = Observer::new(100);
197        for _ in 0..10 { obs.record(EnteSpawned); }
198        for _ in 0..10 { obs.record(EnteDied); }
199        // Bernoulli(0.5) → H = 1 bit.
200        assert!((obs.shannon_entropy() - 1.0).abs() < 1e-9);
201    }
202
203    #[test]
204    fn conditional_prob_perfect_dependency() {
205        let mut obs = Observer::new(100);
206        // Spawned siempre seguido por Died.
207        for _ in 0..5 {
208            obs.record(EnteSpawned);
209            obs.record(EnteDied);
210        }
211        let p = obs.conditional_prob(&EnteSpawned, &EnteDied);
212        assert!(p > 0.0, "esperamos correlación positiva, got {p}");
213    }
214}