arje/crates/ente-brain/src/crystallize.rs

//! Cristalización: del flujo observado a reglas explícitas.
//!
//! Detecta pares (a, b) donde:
//!   - support(a, b) ≥ min_support  (suficientes muestras para no ser ruido)
//!   - P(b|a) ≥ min_conditional_prob (a predice b con confianza)
//!   - PMI(a; b) ≥ min_pmi          (más correlacionados que random)
//!
//! Cada cristal puede emitirse como snippet KCL (texto humano-readable) o
//! como `Rule` ejecutable directamente por el motor.

use crate::observer::Observer;
use crate::rules::{Action, EventKind, EventPattern, LogLevel, Rule, Scope};
use serde::{Deserialize, Serialize};
use ulid::Ulid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Crystal {
    pub antecedent: EventKind,
    pub consequent: EventKind,
    pub conditional_prob: f64,
    pub pmi: f64,
    pub support: u64,
}

#[derive(Debug, Clone, Copy)]
pub struct CrystallizationParams {
    pub min_support: u64,
    pub min_conditional_prob: f64,
    pub min_pmi: f64,
}

impl Default for CrystallizationParams {
    fn default() -> Self {
        Self {
            min_support: 5,
            min_conditional_prob: 0.7,
            min_pmi: 0.5,
        }
    }
}

pub fn detect_crystals(obs: &Observer, params: &CrystallizationParams) -> Vec<Crystal> {
    let mut out = Vec::new();
    for ((a, b), &count) in obs.cooccurrences() {
        if count < params.min_support { continue; }
        let cp = obs.conditional_prob(a, b);
        if cp < params.min_conditional_prob { continue; }
        let mi = obs.pmi(a, b);
        if mi < params.min_pmi { continue; }
        out.push(Crystal {
            antecedent: a.clone(),
            consequent: b.clone(),
            conditional_prob: cp,
            pmi: mi,
            support: count,
        });
    }
    // Orden estable: por confianza descendente para fácil inspección.
    out.sort_by(|x, y| y.conditional_prob.partial_cmp(&x.conditional_prob).unwrap_or(std::cmp::Ordering::Equal));
    out
}

/// Genera un snippet KCL representando la regla cristalizada. El snippet usa
/// la sintaxis tagged union del schema `rule.k` (Single + EventKind nested).
pub fn crystal_to_kcl(c: &Crystal) -> String {
    let id = Ulid::new();
    format!(
r#"# Auto-cristalizado:
#   antecedent → consequent  |  P(c|a) = {cp:.3}, PMI = {pmi:.3} bits, support = {sup}
Rule {{
    id = "{id}"
    priority = 5
    when = EventPattern {{
        type = "Single"
        kind = EventKind {{tag = "{ant_tag}"{ant_extra}}}
    }}
    scope = Scope {{}}
    then = [
        Action {{
            kind = "Log"
            level = "info"
            message = "crystal: {ant_tag} → {con_tag} (auto, P={cp:.2}, PMI={pmi:.2})"
        }}
    ]
}}
"#,
        id = id,
        cp = c.conditional_prob,
        pmi = c.pmi,
        sup = c.support,
        ant_tag = kind_tag(&c.antecedent),
        ant_extra = kind_extra(&c.antecedent),
        con_tag = kind_tag(&c.consequent),
    )
}

fn kind_tag(k: &EventKind) -> &'static str {
    match k {
        EventKind::EnteSpawned => "EnteSpawned",
        EventKind::EnteDied => "EnteDied",
        EventKind::BusAnnounce => "BusAnnounce",
        EventKind::BusInvoke => "BusInvoke",
        EventKind::BusInvokeOf(_) => "BusInvokeOf",
        EventKind::DeviceAdded => "DeviceAdded",
        EventKind::DeviceRemoved => "DeviceRemoved",
        EventKind::Custom(_) => "Custom",
    }
}

fn kind_extra(k: &EventKind) -> String {
    match k {
        EventKind::Custom(s) => format!(", custom = \"{}\"", s.replace('"', "\\\"")),
        // Para BusInvokeOf el cap se omitiría por simplicidad; el snippet
        // promovido es la versión "genérica BusInvoke" salvo que el operador
        // edite manualmente.
        _ => String::new(),
    }
}

/// Convierte un cristal a una `Rule` ejecutable por el motor. Útil para
/// "auto-aprendizaje" donde cristales se promueven a reglas vivas tras
/// validar con el operador.
pub fn crystal_to_rule(c: &Crystal) -> Rule {
    Rule {
        id: Ulid::new(),
        priority: 5,
        when: EventPattern::Single { kind: c.antecedent.clone() },
        scope: Scope::default(),
        then: vec![Action::Log {
            level: LogLevel::Info,
            message: format!(
                "crystal: {:?} → {:?} (P={:.2}, PMI={:.2}, n={})",
                c.antecedent, c.consequent, c.conditional_prob, c.pmi, c.support
            ),
        }],
    }
}


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

    #[test]
    fn detects_perfect_correlation() {
        let mut obs = Observer::new(100);
        for _ in 0..10 {
            obs.record(EnteSpawned);
            obs.record(EnteDied);
        }
        let crystals = detect_crystals(&obs, &CrystallizationParams {
            min_support: 3,
            min_conditional_prob: 0.5,
            min_pmi: 0.0,
        });
        assert!(crystals.iter().any(|c| matches!(c.antecedent, EnteSpawned)
                                       && matches!(c.consequent, EnteDied)));
    }

    #[test]
    fn rejects_below_threshold() {
        let mut obs = Observer::new(100);
        // Sin co-ocurrencia significativa.
        for _ in 0..3 { obs.record(EnteSpawned); }
        let crystals = detect_crystals(&obs, &CrystallizationParams::default());
        assert!(crystals.is_empty(), "no debería haber cristales: {:?}", crystals);
    }
}