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}
30
31impl Observer {
32    pub fn new(window_size: usize) -> Self {
33        Self {
34            window: VecDeque::with_capacity(window_size),
35            window_size,
36            marginal: HashMap::new(),
37            cooccur: HashMap::new(),
38            total: 0,
39        }
40    }
41
42    /// Registra un evento. Actualiza marginales y co-ocurrencias contra todo
43    /// evento aún en la ventana.
44    pub fn record(&mut self, kind: EventKind) {
45        let now = Instant::now();
46        let timed = TimedEvent { kind: kind.clone(), at: now };
47
48        // Co-ocurrencias: este evento con cada uno previo en ventana.
49        for w in &self.window {
50            *self.cooccur
51                .entry((w.kind.clone(), kind.clone()))
52                .or_insert(0) += 1;
53        }
54
55        self.window.push_back(timed);
56        if self.window.len() > self.window_size {
57            self.window.pop_front();
58        }
59
60        *self.marginal.entry(kind).or_insert(0) += 1;
61        self.total += 1;
62    }
63
64    /// Entropía de Shannon de la distribución marginal de eventos.
65    /// H(X) = −Σ p(x) log₂ p(x). Unidad: bits.
66    pub fn shannon_entropy(&self) -> f64 {
67        if self.total == 0 { return 0.0; }
68        let total = self.total as f64;
69        self.marginal.values()
70            .map(|&c| {
71                let p = c as f64 / total;
72                if p > 0.0 { -p * p.log2() } else { 0.0 }
73            })
74            .sum()
75    }
76
77    /// P(b | a) = "dado que algo siguió a `a` dentro del window, qué fracción
78    /// fue `b`". Suma 1 sobre todos los b posibles para un a fijo.
79    ///
80    /// Implementación: cooccur(a, b) / Σ_x cooccur(a, x). Esto da una
81    /// probabilidad condicional propia [0, 1].
82    pub fn conditional_prob(&self, a: &EventKind, b: &EventKind) -> f64 {
83        let joint = self.cooccur
84            .get(&(a.clone(), b.clone()))
85            .copied()
86            .unwrap_or(0) as f64;
87        let row_total: u64 = self.cooccur.iter()
88            .filter_map(|((x, _), c)| if x == a { Some(*c) } else { None })
89            .sum();
90        if row_total == 0 { 0.0 } else { joint / row_total as f64 }
91    }
92
93    /// Información mutua puntual entre `a` y `b`:
94    /// PMI(a, b) = log₂( P(a, b) / (P(a) · P(b)) ).
95    /// Positivo → más correlacionados de lo que sugiere independencia.
96    pub fn pmi(&self, a: &EventKind, b: &EventKind) -> f64 {
97        if self.total == 0 { return 0.0; }
98        let total = self.total as f64;
99        let joint = self.cooccur
100            .get(&(a.clone(), b.clone()))
101            .copied()
102            .unwrap_or(0) as f64 / total;
103        let pa = self.marginal.get(a).copied().unwrap_or(0) as f64 / total;
104        let pb = self.marginal.get(b).copied().unwrap_or(0) as f64 / total;
105        if joint == 0.0 || pa == 0.0 || pb == 0.0 { return 0.0; }
106        (joint / (pa * pb)).log2()
107    }
108
109    /// Información mutua acumulada de la pareja (a, b) ponderada por su
110    /// probabilidad conjunta. Útil como medida de "interés" del par.
111    pub fn weighted_pmi(&self, a: &EventKind, b: &EventKind) -> f64 {
112        if self.total == 0 { return 0.0; }
113        let joint = self.cooccur
114            .get(&(a.clone(), b.clone()))
115            .copied()
116            .unwrap_or(0) as f64 / self.total as f64;
117        joint * self.pmi(a, b)
118    }
119
120    pub fn marginals(&self) -> &HashMap<EventKind, u64> { &self.marginal }
121    pub fn cooccurrences(&self) -> &HashMap<(EventKind, EventKind), u64> { &self.cooccur }
122    pub fn total(&self) -> u64 { self.total }
123    pub fn window_size(&self) -> usize { self.window_size }
124    pub fn current_window(&self) -> usize { self.window.len() }
125
126    /// Últimos N eventos del window, en orden cronológico (más viejo primero).
127    /// Si N > window.len(), devuelve todo el window.
128    pub fn recent(&self, n: usize) -> impl Iterator<Item = &TimedEvent> {
129        let start = self.window.len().saturating_sub(n);
130        self.window.range(start..)
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::rules::EventKind::*;
138
139    #[test]
140    fn entropy_zero_for_single_event() {
141        let mut obs = Observer::new(10);
142        for _ in 0..5 { obs.record(EnteSpawned); }
143        // Distribución degenerada: una sola observación posible → H = 0.
144        assert!(obs.shannon_entropy() < 1e-9);
145    }
146
147    #[test]
148    fn entropy_one_for_balanced_binary() {
149        let mut obs = Observer::new(100);
150        for _ in 0..10 { obs.record(EnteSpawned); }
151        for _ in 0..10 { obs.record(EnteDied); }
152        // Bernoulli(0.5) → H = 1 bit.
153        assert!((obs.shannon_entropy() - 1.0).abs() < 1e-9);
154    }
155
156    #[test]
157    fn conditional_prob_perfect_dependency() {
158        let mut obs = Observer::new(100);
159        // Spawned siempre seguido por Died.
160        for _ in 0..5 {
161            obs.record(EnteSpawned);
162            obs.record(EnteDied);
163        }
164        let p = obs.conditional_prob(&EnteSpawned, &EnteDied);
165        assert!(p > 0.0, "esperamos correlación positiva, got {p}");
166    }
167}