//! CPU-side observer for the GPU-driven train-vs-val state-distribution KL
//! and the derived Flat-trap escape amplifier.
//!
//! Plan 3 Task 7 C.3. The `state_kl_moment_match` GPU kernel reads the
//! train-state sample (`states_out`) and the val-state batch
//! (`chunked_states_buf`), computes per-OFI-dim Gaussian moment-match KL
//! summed over the OFI block, EMAs that into ISV[STATE_KL_TRAIN_VAL_EMA=78],
//! and writes the bounded multiplier ISV[STATE_KL_AMPLIFICATION=79] consumed
//! by `experience_kernels.cu` at the B.1 opp_cost and B.2 bonus sites.
//!
//! This monitor is a pure read-only observer: it surfaces both the EMA and
//! the amp in HEALTH_DIAG and tracks per-epoch fire-rate against the amp so
//! the `controller_activity` smoke test can flag thrash.

use crate::cuda_pipeline::gpu_dqn_trainer::{
    STATE_KL_AMPLIFICATION_INDEX, STATE_KL_TRAIN_VAL_EMA_INDEX,
};
use crate::trainers::dqn::adaptive_monitor::{
    AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus,
};

pub(crate) struct StateKLMonitor {
    /// `None` BEFORE the first observation — encoded as a sentinel so the
    /// first `observe()` call ALWAYS fires (cold-start "we have nothing
    /// to compare against, so call this a fire" semantic). Previously
    /// initialized to `1.0` which coincidentally equals the typical
    /// initial amplification value and silently suppressed the first
    /// fire — the `observe_tracks_fire_rate_on_change` test caught
    /// this regression once exercised.
    last_amp: Option<f32>,
    fire_rate: FireRateStats,
}

impl StateKLMonitor {
    pub(crate) fn new() -> Self {
        Self { last_amp: None, fire_rate: FireRateStats::default() }
    }
}

impl AdaptiveMonitor for StateKLMonitor {
    fn read(&self, isv: &IsvBus<'_>) -> f32 {
        isv.read(STATE_KL_AMPLIFICATION_INDEX)
    }

    fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot {
        let kl_ema = isv.read(STATE_KL_TRAIN_VAL_EMA_INDEX);
        let amp    = isv.read(STATE_KL_AMPLIFICATION_INDEX);
        DiagSnapshot::new()
            .with("state_kl.train_val_ema", kl_ema)
            .with("state_kl.amp",            amp)
            .with("fire_rate",               self.fire_rate.fire_rate())
    }

    fn observe(&mut self, current: f32) {
        let fired = match self.last_amp {
            None => true,                                            // cold-start fire
            Some(prev) => (current - prev).abs() > 1e-7,
        };
        self.fire_rate.record_fire(fired);
        self.last_amp = Some(current);
    }

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

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

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

    fn make_isv(kl_ema: f32, amp: f32) -> Vec<f32> {
        let mut v = vec![0.0f32; STATE_KL_AMPLIFICATION_INDEX + 1];
        v[STATE_KL_TRAIN_VAL_EMA_INDEX] = kl_ema;
        v[STATE_KL_AMPLIFICATION_INDEX] = amp;
        v
    }

    #[test]
    fn read_returns_amp_slot() {
        let mut buf = make_isv(0.42, 1.25);
        let bus = IsvBus::new(&mut buf);
        let mon = StateKLMonitor::new();
        let v = mon.read(&bus);
        assert!((v - 1.25).abs() < 1e-7, "expected amp=1.25, got {v}");
    }

    #[test]
    fn diagnose_exposes_kl_ema_and_amp() {
        let mut buf = make_isv(0.10, 1.50);
        let bus = IsvBus::new(&mut buf);
        let mon = StateKLMonitor::new();
        let snap = mon.diagnose(&bus);
        assert!(snap.fields.contains_key("state_kl.train_val_ema"));
        assert!(snap.fields.contains_key("state_kl.amp"));
        assert!(snap.fields.contains_key("fire_rate"));
        let got_kl  = snap.fields["state_kl.train_val_ema"];
        let got_amp = snap.fields["state_kl.amp"];
        assert!((got_kl - 0.10).abs() < 1e-6);
        assert!((got_amp - 1.50).abs() < 1e-6);
    }

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