//! CPU-side observer for the GPU-driven attention-focus interpretability EMAs
//! introduced in Plan 4 Task 5 Mode A (E.5).
//!
//! The `attention_focus_ema_update` GPU kernel (single-thread cold path)
//! takes 3 host scalars by value once per epoch and EMA-updates:
//!   * ISV[VSN_MAG_EMA_INDEX=87]         — mag-branch VSN-weight magnitude EMA
//!   * ISV[VSN_DIR_EMA_INDEX=88]         — dir-branch equivalent
//!   * ISV[MAMBA2_RETENTION_EMA_INDEX=89] — Mamba2 state-transition magnitude EMA
//!
//! This monitor is a pure read-only observer. It surfaces all three signals in
//! HEALTH_DIAG and tracks fire-rate against the dir-branch slot — that is the
//! signal most likely to thrash if the trunk is unstable, mirroring the
//! convention used by `PlanThresholdMonitor` of picking one anchor slot for
//! fire-rate tracking.
//!
//! Mode A is diagnostic only: no consumer kernel reads slots [87..90).

use crate::cuda_pipeline::gpu_dqn_trainer::{
    MAMBA2_RETENTION_EMA_INDEX, VSN_DIR_EMA_INDEX, VSN_MAG_EMA_INDEX,
};
use crate::trainers::dqn::adaptive_monitor::{
    AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus,
};

pub(crate) struct AttentionMonitor {
    last_dir: f32,
    fire_rate: FireRateStats,
}

impl AttentionMonitor {
    pub(crate) fn new() -> Self {
        Self { last_dir: 0.0, fire_rate: FireRateStats::default() }
    }
}

impl AdaptiveMonitor for AttentionMonitor {
    fn read(&self, isv: &IsvBus<'_>) -> f32 {
        // Anchor read on the direction-branch slot — same convention as
        // PlanThresholdMonitor's `read()` returning the consumer-facing slot.
        isv.read(VSN_DIR_EMA_INDEX)
    }

    fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot {
        let vsn_mag       = isv.read(VSN_MAG_EMA_INDEX);
        let vsn_dir       = isv.read(VSN_DIR_EMA_INDEX);
        let mamba2_retain = isv.read(MAMBA2_RETENTION_EMA_INDEX);
        DiagSnapshot::new()
            .with("attention.vsn_mag",       vsn_mag)
            .with("attention.vsn_dir",       vsn_dir)
            .with("attention.mamba2_retain", mamba2_retain)
            .with("fire_rate",               self.fire_rate.fire_rate())
    }

    fn observe(&mut self, current: f32) {
        let fired = (current - self.last_dir).abs() > 1e-7;
        self.fire_rate.record_fire(fired);
        self.last_dir = current;
    }

    fn fire_rate(&self) -> &FireRateStats {
        &self.fire_rate
    }

    fn name(&self) -> &'static str {
        "attention"
    }
}

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

    fn make_isv(mag: f32, dir: f32, retain: f32) -> Vec<f32> {
        let mut v = vec![0.0f32; MAMBA2_RETENTION_EMA_INDEX + 1];
        v[VSN_MAG_EMA_INDEX]         = mag;
        v[VSN_DIR_EMA_INDEX]         = dir;
        v[MAMBA2_RETENTION_EMA_INDEX] = retain;
        v
    }

    #[test]
    fn read_returns_vsn_dir_anchor_slot() {
        let mut buf = make_isv(0.10, 0.20, 0.30);
        let bus = IsvBus::new(&mut buf);
        let mon = AttentionMonitor::new();
        let v = mon.read(&bus);
        assert!((v - 0.20).abs() < 1e-7, "expected vsn_dir=0.20, got {v}");
    }

    #[test]
    fn diagnose_exposes_all_three_attention_slots() {
        let mut buf = make_isv(0.40, 0.50, 0.60);
        let bus = IsvBus::new(&mut buf);
        let mon = AttentionMonitor::new();
        let snap = mon.diagnose(&bus);
        assert!(snap.fields.contains_key("attention.vsn_mag"));
        assert!(snap.fields.contains_key("attention.vsn_dir"));
        assert!(snap.fields.contains_key("attention.mamba2_retain"));
        assert!(snap.fields.contains_key("fire_rate"));
        assert!((snap.fields["attention.vsn_mag"]       - 0.40).abs() < 1e-6);
        assert!((snap.fields["attention.vsn_dir"]       - 0.50).abs() < 1e-6);
        assert!((snap.fields["attention.mamba2_retain"] - 0.60).abs() < 1e-6);
    }

    #[test]
    fn observe_tracks_fire_rate_on_change() {
        let mut mon = AttentionMonitor::new();
        mon.observe(0.10);
        mon.observe(0.10);
        mon.observe(0.20);
        let rate = mon.fire_rate().fire_rate();
        assert!((rate - 2.0 / 3.0).abs() < 1e-6,
                "expected 2/3 fire rate, got {rate}");
    }
}
