//! N1: Q-network weight snapshots ring buffer for temporal self-distillation.
//!
//! Keeps up to `MAX_SNAPSHOTS` historical weight checkpoints taken at epochs
//! with strong directional Q-differentiation (high q_gap). During collapse
//! recovery, a gradient pull toward the best snapshot is added to the master
//! gradient.
//!
//! Gate design: the snapshot decision is made purely on q_gap quality —
//! q_gap *is* the signal that distillation exists to preserve. Gating on
//! health (the original Layer-1 6-component signal) or winrate (a derived
//! trading metric) was a misstep: the health EMA can be degenerate for
//! small networks, and high-winrate epochs often don't overlap with
//! high-q_gap epochs. Using the direct metric — q_gap — eliminates those
//! indirection layers.

use std::sync::Arc;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
use crate::MLError;

pub const MAX_SNAPSHOTS: usize = 5;
/// Raised from 0.4 → 0.55 after local smoke tests showed 0.4 was too
/// conservative — by the time health drops that low, the Q-collapse
/// attractor has already won. Fire distillation earlier to pull weights
/// back toward the best historical snapshot while there's still structure
/// to preserve.
pub const DISTILL_HEALTH_THRESHOLD: f32 = 0.55;
/// Absolute q_gap sanity floor. Rejects snapshots where q_gap is effectively
/// zero (NaN, denormals, fully-collapsed states). The *quality* floor is
/// dynamic — see `SnapshotRing::dynamic_q_gap_floor` — and scales with the
/// largest q_gap observed in the run. This absolute floor is only a
/// NaN/degenerate guard: a network that has never differentiated any
/// actions at all is not worth distilling toward regardless of scale.
pub const ABSOLUTE_Q_GAP_FLOOR: f32 = 0.001;
/// Fraction of `max_q_gap_observed` required for snapshot acceptance. At
/// 0.5, a new snapshot needs q_gap within 50% of the best q_gap tracked
/// by the decaying peak. Scales automatically with network size: a
/// production-scale network peaking at 0.5 gets a floor of 0.25 (strict);
/// a smoke-test network peaking at 0.05 gets a floor of 0.025 (permissive).
pub const DYNAMIC_Q_GAP_FLOOR_FRACTION: f32 = 0.5;
/// Per-epoch decay factor for `max_q_gap_observed`. The peak tracker uses
/// `peak = max(peak * decay, q_gap)` — rises instantly to a new peak,
/// relaxes slowly. 0.99/epoch yields a ~70-epoch half-life: a lucky early
/// peak doesn't permanently lock out later snapshots, and a regime shift
/// that reduces the natural Q-gap scale is eventually accommodated. Plain
/// running-max (no decay) would trap the floor at any outlier. Plain EMA
/// (of q_gap itself) would drop during collapse exactly when the floor
/// needs to stay up. Decaying peak is the right asymmetry for rollback
/// high-water marks — same pattern used by backtracking checkpoints.
pub const PEAK_DECAY_PER_EPOCH: f32 = 0.99;

/// A single weight snapshot with metadata.
pub struct QSnapshot {
    pub weights: CudaSlice<f32>,
    pub health: f32,
    pub q_gap: f32,
    pub epoch: u32,
}

impl std::fmt::Debug for QSnapshot {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("QSnapshot")
            .field("health", &self.health)
            .field("q_gap", &self.q_gap)
            .field("epoch", &self.epoch)
            .finish()
    }
}

/// Ring buffer of best-health snapshots.
pub struct SnapshotRing {
    stream: Arc<CudaStream>,
    pub(crate) snapshots: Vec<QSnapshot>,
    param_count: usize,
    /// Running maximum q_gap observed across the run. Updated on every
    /// `maybe_snapshot` call (accepts and rejects both). Drives the dynamic
    /// q_gap floor so the snapshot gate adapts to the network's natural
    /// Q-gap magnitude instead of relying on a scale-specific constant.
    max_q_gap_observed: f32,
}

impl std::fmt::Debug for SnapshotRing {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SnapshotRing")
            .field("param_count", &self.param_count)
            .field("snapshot_count", &self.snapshots.len())
            .finish()
    }
}

impl SnapshotRing {
    pub fn new(stream: Arc<CudaStream>, param_count: usize) -> Self {
        Self {
            stream,
            snapshots: Vec::with_capacity(MAX_SNAPSHOTS),
            param_count,
            max_q_gap_observed: 0.0,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.snapshots.is_empty()
    }

    pub fn param_count(&self) -> usize {
        self.param_count
    }

    /// Current dynamic q_gap acceptance floor. Scales with the largest q_gap
    /// seen in the run, clamped below by `ABSOLUTE_Q_GAP_FLOOR` so degenerate
    /// runs (q_gap never exceeds zero) still reject NaN/zero snapshots.
    pub fn dynamic_q_gap_floor(&self) -> f32 {
        (DYNAMIC_Q_GAP_FLOOR_FRACTION * self.max_q_gap_observed).max(ABSOLUTE_Q_GAP_FLOOR)
    }

    /// Register an observed q_gap, updating the decaying-peak tracker.
    /// Semantics: `peak = max(peak * decay, q_gap)` — rises instantly to a
    /// new peak, relaxes by `PEAK_DECAY_PER_EPOCH` otherwise. Callers should
    /// invoke this once per epoch boundary on every snapshot attempt (accept
    /// and reject both) so the floor reflects the true per-run q_gap
    /// distribution and adapts to regime shifts over long horizons.
    pub fn observe_q_gap(&mut self, q_gap: f32) {
        let decayed = self.max_q_gap_observed * PEAK_DECAY_PER_EPOCH;
        let observed = if q_gap.is_finite() { q_gap } else { 0.0 };
        self.max_q_gap_observed = decayed.max(observed);
    }

    /// Copy current weights into a new snapshot if q_gap is above the
    /// dynamic floor AND beats the worst stored entry (when at capacity).
    ///
    /// The `health` parameter is stored on the snapshot as metadata but
    /// does not gate acceptance — q_gap is the direct measure of the thing
    /// we want to preserve (directional Q-differentiation), and layering
    /// a multi-component health proxy on top of it cost us captures in
    /// runs where health's internal EMAs were degenerate despite strong
    /// q_gap peaks.
    pub fn maybe_snapshot(
        &mut self,
        params_src: &CudaSlice<f32>,
        health: f32,
        q_gap: f32,
        epoch: u32,
    ) -> Result<bool, MLError> {
        // Track every observation (accept or reject) so the dynamic floor
        // reflects the true per-run q_gap distribution, including peaks
        // at epochs that the floor rejects below.
        self.observe_q_gap(q_gap);

        // Dynamic q_gap floor: scales with the decaying peak observed in
        // this run. Absolute floor of 0.001 catches NaN / fully-degenerate
        // networks; the relative floor (0.5 × decaying peak) catches any
        // state meaningfully worse than the best this run has seen.
        if q_gap < self.dynamic_q_gap_floor() {
            return Ok(false);
        }

        // At capacity, only replace if this q_gap beats the worst stored entry.
        if self.snapshots.len() >= MAX_SNAPSHOTS {
            let worst_q_gap = self.snapshots.iter()
                .map(|s| s.q_gap)
                .fold(f32::INFINITY, f32::min);
            if q_gap <= worst_q_gap {
                return Ok(false);
            }
        }

        let n = self.param_count;
        let num_bytes = n * std::mem::size_of::<f32>();
        // SAFETY: `alloc` returns uninitialized GPU memory, which we immediately
        // overwrite with `memcpy_dtod` below before any read. No UB window.
        #[allow(unsafe_code)]
        let mut new_weights = unsafe { self.stream.alloc::<f32>(n) }
            .map_err(|e| MLError::ModelError(format!("snapshot alloc: {e}")))?;
        {
            let (src_ptr, _src_guard) = params_src.device_ptr(&self.stream);
            let (dst_ptr, _dst_guard) = new_weights.device_ptr_mut(&self.stream);
            #[allow(unsafe_code)]
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(),
                ).map_err(|e| MLError::ModelError(format!("snapshot dtod: {e}")))?;
            }
        }

        if self.snapshots.len() >= MAX_SNAPSHOTS {
            let worst_idx = self.snapshots.iter()
                .enumerate()
                .min_by(|(_, a), (_, b)| {
                    // NaN-safe: treat NaN q_gaps as equal so a snapshot with NaN
                    // does not panic the swap-out path.
                    a.q_gap.partial_cmp(&b.q_gap).unwrap_or(std::cmp::Ordering::Equal)
                })
                .map(|(i, _)| i)
                .expect("MAX_SNAPSHOTS >= 1 and len >= MAX_SNAPSHOTS, so at least one snapshot exists");
            self.snapshots.swap_remove(worst_idx);
        }

        self.snapshots.push(QSnapshot { weights: new_weights, health, q_gap, epoch });
        Ok(true)
    }

    /// Best (highest q_gap) snapshot raw device pointer, if any.
    pub fn best_raw_ptr(&self) -> Option<u64> {
        self.snapshots.iter()
            .max_by(|a, b| a.q_gap.partial_cmp(&b.q_gap).unwrap_or(std::cmp::Ordering::Equal))
            .map(|s| s.weights.raw_ptr())
    }
}
