//! Phase R5 gate G4: `DqnHead::soft_update_target` implements the
//! Polyak averaging formula
//!
//!   target[i] = (1 − τ)·target[i] + τ·current[i]
//!
//! with τ read from `ISV[RL_TARGET_TAU_INDEX]`.
//!
//! Per `feedback_no_cpu_test_fallbacks`: the oracle is the algebraic
//! identity the kernel implements, evaluated host-side on the same
//! numbers the kernel saw — not a CPU reimplementation.
//!
//! G3 (which exercised the now-removed `launch_rl_controllers_per_step`
//! bulk method) was retired when the trainer adopted
//! `launch_rl_fused_controllers` — a single fused kernel reading from
//! per-step dones + a dedicated input-slot index buffer. The fused
//! kernel's "all controllers move slots" invariant is exercised end-to-end
//! by every cluster training run, so a separate unit test would
//! reproduce production setup without adding signal.
//!
//! Run with:
//!   `cargo test -p ml-alpha --test r5_controllers_and_soft_update -- --ignored --nocapture`

use ml_alpha::rl::isv_slots::RL_TARGET_TAU_INDEX;
use ml_alpha::trainer::integrated::{IntegratedTrainer, IntegratedTrainerConfig};
use ml_alpha::trainer::perception::PerceptionTrainerConfig;
use ml_core::device::MlDevice;

/// Bootstrap value the `rl_target_tau_controller` writes into
/// `ISV[RL_TARGET_TAU_INDEX]` at construction time (sentinel-input
/// path of `pearl_first_observation_bootstrap`).
const TAU_BOOTSTRAP: f32 = 0.005;

fn build_trainer() -> Option<(MlDevice, IntegratedTrainer)> {
    let dev = match MlDevice::cuda(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("CUDA 0 not available — skipping ({e})");
            return None;
        }
    };
    let cfg = IntegratedTrainerConfig {
        perception: PerceptionTrainerConfig {
            seq_len: 4,
            n_batch: 1,
            ..PerceptionTrainerConfig::default()
        },
        dqn_seed: 0xB7,
        ppo_seed: 0xB8,
        ..IntegratedTrainerConfig::default()
    };
    let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new");
    Some((dev, trainer))
}

#[test]
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
fn g4_dqn_target_soft_update_implements_polyak_formula() {
    let Some((dev, mut trainer)) = build_trainer() else { return };
    let stream = dev.cuda_stream().expect("cuda_stream").clone();

    // Force-overwrite w_d with all-ones to break the
    // (w_d == w_target_d) symmetry the DqnHead init establishes.
    // (Otherwise the soft update has nothing to blend toward.)
    let n_w = trainer.dqn_head.w_d.len();
    let ones = vec![1.0_f32; n_w];
    stream
        .memcpy_htod(&ones, &mut trainer.dqn_head.w_d)
        .expect("force w_d");

    // Snapshot w_target BEFORE soft_update. This is the Xavier init
    // (small random values).
    let mut target_before = vec![0.0_f32; n_w];
    stream
        .memcpy_dtoh(&trainer.dqn_head.w_target_d, target_before.as_mut_slice())
        .expect("dtoh w_target before");

    // τ = ISV[RL_TARGET_TAU_INDEX] bootstrap value. Sync the trainer's
    // stream first so the bootstrap-controller writes from
    // `IntegratedTrainer::new` are visible through the mapped-pinned
    // ISV host view per `feedback_no_htod_htoh_only_mapped_pinned`.
    trainer.stream.synchronize().expect("sync trainer stream");
    let tau = trainer.read_isv_host(RL_TARGET_TAU_INDEX);
    assert!(
        (tau - TAU_BOOTSTRAP).abs() < 1e-6,
        "pre-condition: τ should be R1-bootstrapped to {TAU_BOOTSTRAP}; got {tau}"
    );

    // Fire the soft update. `soft_update_target` reads τ from
    // `isv_dev_ptr` (raw `CUdeviceptr`, zero-copy stable pointer).
    let isv_ptr = trainer.isv_dev_ptr;
    trainer
        .dqn_head
        .soft_update_target(&isv_ptr)
        .expect("soft_update_target");
    stream.synchronize().expect("sync after soft_update");

    // Snapshot w_target AFTER. For each element:
    //   target_after[i] = (1 − τ) · target_before[i] + τ · w_d[i]
    //                   = 0.995 · target_before[i] + 0.005 · 1.0
    let mut target_after = vec![0.0_f32; n_w];
    stream
        .memcpy_dtoh(&trainer.dqn_head.w_target_d, target_after.as_mut_slice())
        .expect("dtoh w_target after");

    // Spot-check first / mid / last indices.
    let sample_indices = [0_usize, n_w / 2, n_w - 1];
    for &i in &sample_indices {
        let expected = (1.0 - tau) * target_before[i] + tau * 1.0;
        let got = target_after[i];
        assert!(
            (got - expected).abs() < 1e-6,
            "soft_update target[{i}]: expected (1-τ)·{}+τ·1 = {expected}, got {got}",
            target_before[i]
        );
    }

    // Negative invariant: at least one element must have CHANGED
    // (the formula moves every element by `τ · (w_d[i] - target[i])`,
    // which is non-zero unless w_d[i] == target[i] for all i — false
    // here since w_d is all-ones and target_before is Xavier-random).
    let any_changed = (0..n_w).any(|i| (target_after[i] - target_before[i]).abs() > 1e-9);
    assert!(
        any_changed,
        "soft_update should change at least one target element when w_d != target"
    );

    eprintln!(
        "G4 OK — soft_update applied Polyak formula with τ={tau}: \
         target[0] {} → {} (expected {})",
        target_before[0],
        target_after[0],
        (1.0 - tau) * target_before[0] + tau * 1.0,
    );
}
