//! CPU-side observer for the GPU-driven per-component reward EMAs.
//!
//! The `reward_component_ema` CUDA kernel reads `reward_components_per_sample`
//! [N*L, 6] written by `experience_env_step` and updates 6 ISV EMA slots
//! [REWARD_POPART_EMA_INDEX..REWARD_POPART_EMA_INDEX+6) = [63..69).
//! This monitor reads those slots for HEALTH_DIAG emission and fire-rate tracking.
//!
//! The monitor's `read()` returns the mean |popart reward| (ISV[63]) as the
//! primary scalar (the largest-magnitude component, representing overall reward
//! scale). `diagnose()` exposes all six components plus their mean and fire_rate.

use crate::cuda_pipeline::gpu_dqn_trainer::{
    REWARD_POPART_EMA_INDEX, REWARD_CF_EMA_INDEX, REWARD_TRAIL_EMA_INDEX,
    REWARD_MICRO_EMA_INDEX, REWARD_OPP_COST_EMA_INDEX, REWARD_BONUS_EMA_INDEX,
};
use crate::trainers::dqn::adaptive_monitor::{
    AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus,
};

pub(crate) struct RewardComponentMonitor {
    last_popart: f32,
    fire_rate: FireRateStats,
}

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

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

    fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot {
        let popart   = isv.read(REWARD_POPART_EMA_INDEX);
        let cf       = isv.read(REWARD_CF_EMA_INDEX);
        let trail    = isv.read(REWARD_TRAIL_EMA_INDEX);
        let micro    = isv.read(REWARD_MICRO_EMA_INDEX);
        let opp_cost = isv.read(REWARD_OPP_COST_EMA_INDEX);
        let bonus    = isv.read(REWARD_BONUS_EMA_INDEX);
        let mean     = (popart + cf + trail + micro + opp_cost + bonus) / 6.0;
        DiagSnapshot::new()
            .with("popart",   popart)
            .with("cf",       cf)
            .with("trail",    trail)
            .with("micro",    micro)
            .with("opp_cost", opp_cost)
            .with("bonus",    bonus)
            .with("mean",     mean)
            .with("fire_rate", self.fire_rate.fire_rate())
    }

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

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

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

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

    fn make_isv(popart: f32, cf: f32, trail: f32, micro: f32, opp_cost: f32, bonus: f32) -> Vec<f32> {
        let mut v = vec![0.0f32; REWARD_BONUS_EMA_INDEX + 1];
        v[REWARD_POPART_EMA_INDEX]   = popart;
        v[REWARD_CF_EMA_INDEX]       = cf;
        v[REWARD_TRAIL_EMA_INDEX]    = trail;
        v[REWARD_MICRO_EMA_INDEX]    = micro;
        v[REWARD_OPP_COST_EMA_INDEX] = opp_cost;
        v[REWARD_BONUS_EMA_INDEX]    = bonus;
        v
    }

    #[test]
    fn read_returns_popart_ema_slot() {
        let mut buf = make_isv(0.25, 0.10, 0.0, 0.05, 0.0, 0.0);
        let bus = IsvBus::new(&mut buf);
        let mon = RewardComponentMonitor::new();
        let v = mon.read(&bus);
        assert!((v - 0.25).abs() < 1e-7, "expected popart=0.25, got {v}");
    }

    #[test]
    fn diagnose_exposes_all_six_components_and_mean() {
        let mut buf = make_isv(0.20, 0.10, 0.0, 0.05, 0.0, 0.0);
        let bus = IsvBus::new(&mut buf);
        let mon = RewardComponentMonitor::new();
        let snap = mon.diagnose(&bus);
        assert!(snap.fields.contains_key("popart"));
        assert!(snap.fields.contains_key("cf"));
        assert!(snap.fields.contains_key("trail"));
        assert!(snap.fields.contains_key("micro"));
        assert!(snap.fields.contains_key("opp_cost"));
        assert!(snap.fields.contains_key("bonus"));
        assert!(snap.fields.contains_key("mean"));
        let expected_mean = (0.20 + 0.10 + 0.0 + 0.05 + 0.0 + 0.0) / 6.0;
        let got_mean = snap.fields["mean"];
        assert!((got_mean - expected_mean).abs() < 1e-6,
                "expected mean={expected_mean:.6}, got {got_mean:.6}");
    }

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