//! Phase R3 gate — three GPU-oracle tests for the new generic kernels:
//!
//!   1. `ema_update_on_done` — bootstrap path: sentinel-zero ISV slot
//!      + first observation k → ISV[slot] == k exactly.
//!   2. `ema_update_per_step` — convergence: constant input k for many
//!      steps → ISV[slot] → k within tolerance.
//!   3. `compute_advantage_return` — formula: zero rewards + done=0 +
//!      `v_t = v_tp1 = k` + γ=0.99 (from R1 bootstrap) → `returns = 0.99k`,
//!      `advantages = (0.99 − 1)·k = −0.01k`.
//!
//! Per `feedback_no_cpu_test_fallbacks` every oracle here is either:
//!   - the kernel's documented bootstrap behaviour (kernel writes
//!     `mean_obs` directly on first observation), OR
//!   - an analytical property of the formula (EMA of constant =
//!     constant; advantage = γ·V − V = (γ−1)·V when r=0, done=0,
//!     v_t = v_tp1), OR
//!   - the kernel's own `*_BOOTSTRAP` value (γ=0.99 from R1).
//!
//! Per `pearl_tests_must_prove_not_lock_observations` each assert is
//! an invariant statement, not a tuned magic number.
//!
//! Run with:
//!   `cargo test -p ml-alpha --test r3_ema_advantage -- --ignored --nocapture`

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

const ALPHA_FLOOR: f32 = 0.4;

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: 0xA1,
        ppo_seed: 0xA2,
        ..IntegratedTrainerConfig::default()
    };
    let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new");
    Some((dev, trainer))
}

fn upload(
    stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    host: &[f32],
) -> cudarc::driver::CudaSlice<f32> {
    let mut d = stream.alloc_zeros::<f32>(host.len()).expect("alloc");
    stream.memcpy_htod(host, &mut d).expect("htod");
    d
}

/// Mapped-pinned ISV readback per `feedback_no_htod_htoh_only_mapped_pinned`:
/// the trainer's `isv_mapped` is the same buffer the GPU writes via
/// `isv_dev_ptr`, so the host view is zero-copy. Caller must have
/// synchronized the producing stream first.
fn readback_isv(trainer: &ml_alpha::trainer::integrated::IntegratedTrainer) -> Vec<f32> {
    trainer.isv_host_slice().to_vec()
}

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

    // Pre-condition: the EMA-input slot is at sentinel zero (R1
    // bootstraps ISV[400..406] for controllers; ISV[417..423] EMA-input
    // slots stay at alloc_zeros per the R1 invariant).
    let isv_before = readback_isv(&trainer);
    assert_eq!(
        isv_before[RL_MEAN_ABS_PNL_EMA_INDEX], 0.0,
        "pre-condition: EMA slot must be sentinel zero"
    );

    // One observation: obs = [7.0], dones = [1.0]. Per
    // pearl_first_observation_bootstrap the kernel writes the
    // observation directly into the slot — no blend with the cold-
    // start zero.
    let obs_d = upload(&stream, &[7.0]);
    let dones_d = upload(&stream, &[1.0]);
    trainer
        .launch_ema_update_on_done(RL_MEAN_ABS_PNL_EMA_INDEX, ALPHA_FLOOR, &obs_d, &dones_d, 1)
        .expect("ema_update_on_done");
    stream.synchronize().expect("sync");

    let isv_after = readback_isv(&trainer);
    let val = isv_after[RL_MEAN_ABS_PNL_EMA_INDEX];
    // Exact equality — bootstrap path is `isv[slot] = mean_obs` with
    // no arithmetic. Any drift indicates a wrong code path.
    assert_eq!(
        val, 7.0,
        "first-observation bootstrap should write mean_obs exactly; got {val}"
    );

    // Negative invariant: hold-only step (no dones) MUST NOT touch
    // the slot — defends against an off-by-one in the gate count.
    let obs2_d = upload(&stream, &[999.0]);
    let dones2_d = upload(&stream, &[0.0]);
    trainer
        .launch_ema_update_on_done(RL_MEAN_ABS_PNL_EMA_INDEX, ALPHA_FLOOR, &obs2_d, &dones2_d, 1)
        .expect("ema_update_on_done hold");
    stream.synchronize().expect("sync");

    let isv_hold = readback_isv(&trainer);
    assert_eq!(
        isv_hold[RL_MEAN_ABS_PNL_EMA_INDEX], 7.0,
        "hold step (no done) must preserve EMA, not blend toward 0"
    );

    eprintln!(
        "R3.1 OK — bootstrap wrote 7.0 directly; hold-only step preserved value"
    );
}

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

    // Property: EMA of constant k converges to k. With α=0.4 and
    // bootstrap-then-blend, after N=50 steps with obs=5.0:
    //   step 0 (bootstrap): isv = 5.0 directly
    //   step 1+ (blend):    isv = (1-α)·prev + α·5.0 = stays at 5.0
    // So the assertion is essentially exact (modulo fp rounding).
    let k = 5.0_f32;
    let obs_d = upload(&stream, &[k]);
    for _ in 0..50 {
        trainer
            .launch_ema_update_per_step(RL_KL_PI_EMA_INDEX, ALPHA_FLOOR, &obs_d, 1)
            .expect("ema_update_per_step");
    }
    stream.synchronize().expect("sync");

    let isv = readback_isv(&trainer);
    let val = isv[RL_KL_PI_EMA_INDEX];
    assert!(
        (val - k).abs() < 1e-4,
        "EMA of constant {k} should converge to {k}; got {val}"
    );

    eprintln!("R3.2 OK — EMA(constant={k}) converged to {val} after 50 steps");
}

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

    // R1 bootstrapped ISV[γ=400] via rl_gamma_controller. Post-R9
    // audit, the bootstrap derives from the input EMA via the same
    // target formula the per-step path uses — at sentinel input it
    // lands on the floor γ = GAMMA_MIN = 0.90. The test below
    // computes its expected values from whatever γ ISV holds, so the
    // pre-condition is just "γ is in the valid bounded range" rather
    // than a hardcoded canonical value.
    let isv = readback_isv(&trainer);
    let gamma = isv[RL_GAMMA_INDEX];
    assert!(
        gamma >= 0.90 && gamma <= 0.999,
        "pre-condition: γ should be in [0.90, 0.999]; got {gamma}"
    );

    // Inputs: r=0, done=0, v_t=v_tp1=k. Expected:
    //   returns[b]    = 0 + γ·1·k = γ·k
    //   advantages[b] = γ·k − k   = (γ − 1)·k
    let k = 5.0_f32;
    let b_size = 3;
    let rewards_d = upload(&stream, &vec![0.0_f32; b_size]);
    let dones_d = upload(&stream, &vec![0.0_f32; b_size]);
    let v_t_d = upload(&stream, &vec![k; b_size]);
    let v_tp1_d = upload(&stream, &vec![k; b_size]);
    let mut returns_d = stream.alloc_zeros::<f32>(b_size).expect("alloc returns");
    let mut advantages_d = stream.alloc_zeros::<f32>(b_size).expect("alloc advantages");

    trainer
        .launch_compute_advantage_return(
            &rewards_d,
            &dones_d,
            &v_t_d,
            &v_tp1_d,
            &mut returns_d,
            &mut advantages_d,
            b_size,
        )
        .expect("compute_advantage_return");

    let mut returns = vec![0.0_f32; b_size];
    let mut advantages = vec![0.0_f32; b_size];
    stream.memcpy_dtoh(&returns_d, returns.as_mut_slice()).expect("dtoh returns");
    stream.memcpy_dtoh(&advantages_d, advantages.as_mut_slice()).expect("dtoh advantages");
    stream.synchronize().expect("sync");

    let expected_ret = gamma * k; // 0.99 · 5 = 4.95
    let expected_adv = (gamma - 1.0) * k; // -0.05
    for b in 0..b_size {
        assert!(
            (returns[b] - expected_ret).abs() < 1e-5,
            "returns[{b}] expected {expected_ret}, got {}",
            returns[b]
        );
        assert!(
            (advantages[b] - expected_adv).abs() < 1e-5,
            "advantages[{b}] expected {expected_adv}, got {}",
            advantages[b]
        );
    }

    // Negative invariant: done=1 zeros the future-value bootstrap.
    let dones_done_d = upload(&stream, &vec![1.0_f32; b_size]);
    trainer
        .launch_compute_advantage_return(
            &rewards_d,
            &dones_done_d,
            &v_t_d,
            &v_tp1_d,
            &mut returns_d,
            &mut advantages_d,
            b_size,
        )
        .expect("compute_advantage_return done=1");
    stream.memcpy_dtoh(&returns_d, returns.as_mut_slice()).expect("dtoh returns done");
    stream.memcpy_dtoh(&advantages_d, advantages.as_mut_slice()).expect("dtoh advantages done");
    stream.synchronize().expect("sync");

    for b in 0..b_size {
        assert!(
            returns[b].abs() < 1e-6,
            "returns[{b}] with done=1 + r=0 expected 0, got {}",
            returns[b]
        );
        assert!(
            (advantages[b] - (-k)).abs() < 1e-5,
            "advantages[{b}] with done=1 expected −k = {}, got {}",
            -k,
            advantages[b]
        );
    }

    eprintln!(
        "R3.3 OK — γ={gamma}: returns=4.95 / advantages=-0.05 when v=5; \
         done=1 zeros bootstrap → returns=0 / advantages=-5"
    );
}
