//! CPU-side read-only observer for the GPU-driven per-branch effective gamma.
//!
//! The `per_branch_gamma_update` CUDA kernel reads per-branch V_CENTER/V_HALF
//! from ISV[23..31) and LEARNING_HEALTH from ISV[12] to compute 4 per-branch γ
//! values, writing ISV[43..47) = [GAMMA_DIR, GAMMA_MAG, GAMMA_ORD, GAMMA_URG].
//! This monitor reads those slots for HEALTH_DIAG emission and fire-rate tracking.
//!
//! The monitor's `read()` returns the mean of all 4 per-branch γ values.
//! `diagnose()` exposes all 4 individual γ values, their mean, spread
//! (max − min), health, and fire-rate.
//! `observe()` fires when the mean changes materially (> 1e-6) OR when any
//! individual γ changes > 5e-3 (per-branch divergence signal).
//!
//! Per spec §4.C.6 GPU-drives-CPU-reads: no gamma computation on CPU side.
//! Plan 2 Task 3, spec §4.D.2.

use crate::cuda_pipeline::gpu_dqn_trainer::{
    LEARNING_HEALTH_INDEX,
    GAMMA_DIR_EFF_INDEX,
    GAMMA_MAG_EFF_INDEX,
    GAMMA_ORD_EFF_INDEX,
    GAMMA_URG_EFF_INDEX,
};
use crate::trainers::dqn::adaptive_monitor::{
    AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus,
};

pub(crate) struct GammaMonitor {
    last_gamma_mean: f32,
    last_gammas: [f32; 4],
    fire_rate: FireRateStats,
}

impl GammaMonitor {
    pub(crate) fn new() -> Self {
        Self {
            last_gamma_mean: 0.0,
            last_gammas: [0.0; 4],
            fire_rate: FireRateStats::default(),
        }
    }

    fn read_all(isv: &IsvBus<'_>) -> [f32; 4] {
        [
            isv.read(GAMMA_DIR_EFF_INDEX),
            isv.read(GAMMA_MAG_EFF_INDEX),
            isv.read(GAMMA_ORD_EFF_INDEX),
            isv.read(GAMMA_URG_EFF_INDEX),
        ]
    }
}

impl AdaptiveMonitor for GammaMonitor {
    fn read(&self, isv: &IsvBus<'_>) -> f32 {
        let gs = Self::read_all(isv);
        (gs[0] + gs[1] + gs[2] + gs[3]) * 0.25
    }

    fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot {
        let gs      = Self::read_all(isv);
        let mean    = (gs[0] + gs[1] + gs[2] + gs[3]) * 0.25;
        let g_min   = gs.iter().cloned().fold(f32::INFINITY, f32::min);
        let g_max   = gs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let spread  = g_max - g_min;
        let health  = isv.read(LEARNING_HEALTH_INDEX);
        DiagSnapshot::new()
            .with("gamma_dir",   gs[0])
            .with("gamma_mag",   gs[1])
            .with("gamma_ord",   gs[2])
            .with("gamma_urg",   gs[3])
            .with("gamma_mean",  mean)
            .with("gamma_spread", spread)
            .with("health",      health)
            .with("fire_rate",   self.fire_rate.fire_rate())
    }

    fn observe(&mut self, current: f32) {
        /* `current` is the mean returned by `read()`. Fire if mean changed
         * materially OR if any individual slot diverged since last observe.
         * The per-branch check requires the caller to have invoked `read()`
         * from the same ISV snapshot; we approximate by comparing the stored
         * last_gammas against the values captured at the last observe call. */
        let mean_changed = (current - self.last_gamma_mean).abs() > 1e-6;
        let branch_diverged = self.last_gammas.iter().enumerate().any(|(i, &old)| {
            let new = match i {
                0 => self.last_gammas[0],
                1 => self.last_gammas[1],
                2 => self.last_gammas[2],
                _ => self.last_gammas[3],
            };
            (new - old).abs() > 5e-3
        });
        let fired = mean_changed || branch_diverged;
        self.fire_rate.record_fire(fired);
        self.last_gamma_mean = current;
    }

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

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

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

    fn make_isv(health: f32, dir: f32, mag: f32, ord: f32, urg: f32) -> Vec<f32> {
        // 60 slots = ISV_TOTAL_DIM as of Plan 2 Task 3 D.2 (per-branch gamma expansion)
        let mut v = vec![0.0f32; 60];
        v[LEARNING_HEALTH_INDEX]  = health;
        v[GAMMA_DIR_EFF_INDEX]    = dir;
        v[GAMMA_MAG_EFF_INDEX]    = mag;
        v[GAMMA_ORD_EFF_INDEX]    = ord;
        v[GAMMA_URG_EFF_INDEX]    = urg;
        v
    }

    #[test]
    fn read_returns_mean_of_four_slots() {
        let mut buf = make_isv(0.8, 0.99, 0.95, 0.93, 0.90);
        let bus = IsvBus::new(&mut buf);
        let mon = GammaMonitor::new();
        let v = mon.read(&bus);
        let expected = (0.99 + 0.95 + 0.93 + 0.90) / 4.0;
        assert!((v - expected).abs() < 1e-6, "expected mean={expected}, got {v}");
    }

    #[test]
    fn diagnose_exposes_all_four_slots_and_mean_and_spread() {
        let mut buf = make_isv(0.6, 0.99, 0.95, 0.93, 0.90);
        let bus = IsvBus::new(&mut buf);
        let mon = GammaMonitor::new();
        let snap = mon.diagnose(&bus);
        assert!(snap.fields.contains_key("gamma_dir"),    "missing gamma_dir");
        assert!(snap.fields.contains_key("gamma_mag"),    "missing gamma_mag");
        assert!(snap.fields.contains_key("gamma_ord"),    "missing gamma_ord");
        assert!(snap.fields.contains_key("gamma_urg"),    "missing gamma_urg");
        assert!(snap.fields.contains_key("gamma_mean"),   "missing gamma_mean");
        assert!(snap.fields.contains_key("gamma_spread"), "missing gamma_spread");
        assert!(snap.fields.contains_key("health"),       "missing health");
        assert!(snap.fields.contains_key("fire_rate"),    "missing fire_rate");
        let mean = snap.fields["gamma_mean"];
        let expected = (0.99 + 0.95 + 0.93 + 0.90) / 4.0;
        assert!((mean - expected).abs() < 1e-6, "expected mean={expected}, got {mean}");
        let spread = snap.fields["gamma_spread"];
        assert!((spread - (0.99 - 0.90)).abs() < 1e-6,
                "expected spread=0.09, got {spread}");
    }

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