//! Mechanical debug of PerceptionTrainer's grad + weight movement.
//!
//! Prints grad magnitudes, weight deltas, and probs across the first
//! 5 training steps on a single fixed direction=+1 input. The signal
//! is constant (label always 1), so ANY working gradient chain should
//! drive probs from 0.5 upward within 5 steps. Failure modes this
//! catches:
//!   - grad_heads_w / grad_heads_b near zero (broken bwd chain)
//!   - weights not changing after AdamW.step() (silent no-op)
//!   - probs unchanged across steps (forward not using updated weights)

use cudarc::driver::{CudaSlice, CudaStream, DevicePtr};
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::pinned_mem::MappedF32Buffer;
use ml_alpha::trainer::perception::{PerceptionTrainer, PerceptionTrainerConfig};
use ml_core::device::MlDevice;
use std::sync::Arc;

fn test_device() -> MlDevice {
    MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests")
}

fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Vec<f32> {
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }.unwrap();
    let nbytes = n * 4;
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()).unwrap();
    }
    stream.synchronize().unwrap();
    staging.read_all()
}

fn fixed_up_input(prev_mid: f32, ts_ns: u64, prev_ts_ns: u64) -> (Mbp10RawInput, [f32; 5]) {
    let next_mid = prev_mid + 0.25;
    let mut bid_px = [0.0; 10];
    let mut bid_sz = [0.0; 10];
    let mut ask_px = [0.0; 10];
    let mut ask_sz = [0.0; 10];
    for i in 0..10 {
        bid_px[i] = next_mid - 0.125 - 0.25 * i as f32;
        ask_px[i] = next_mid + 0.125 + 0.25 * i as f32;
        bid_sz[i] = 10.0;
        ask_sz[i] = 10.0;
    }
    (
        Mbp10RawInput {
            bid_px, bid_sz, ask_px, ask_sz,
            prev_mid,
            trade_signed_vol: 1.0,
            trade_count: 1,
            ts_ns,
            prev_ts_ns,
        },
        [1.0, 1.0, 1.0, 1.0, 1.0],
    )
}

#[test]
fn debug_long_horizon_weight_trajectory() {
    // Same trainer + constant +1 input, but tracks key state over
    // 200 steps. If 5-step run shows loss DROPPING but 200-step run
    // ends at loss ~8 (probs at WRONG tail), bisect: when does the
    // sign actually flip? Prints every 10 steps.
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        n_in: 32,
        n_hid: 32,
        seq_len: 1,
        lr: 3e-2,
        weight_decay: 0.0,
        seed: 0x4242,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let stream = dev.cuda_stream().unwrap().clone();

    let mut ts_ns = 1_000_000u64;
    let mut prev_ts_ns = 0u64;
    let mut prev_mid = 5500.0_f32;
    println!("step | loss | probs[0] | heads_b[0] | heads_w[0,0] | b_cfc[0]");
    for step in 0..200 {
        trainer.reset_hidden_state().expect("reset");
        let (input, labels) = fixed_up_input(prev_mid, ts_ns, prev_ts_ns);
        let loss = trainer.step(&input, &labels).expect("step");
        if step % 10 == 0 || step < 5 {
            // Need a probs read AFTER training — re-run forward only.
            // Easier: use heads_b movement as a sentinel.
            let heads_b = download(&stream, &trainer.heads_b_d);
            let heads_w = download(&stream, &trainer.heads_w_d);
            let b_cfc = download(&stream, &trainer.b_d);
            println!(
                "{step:4} | {loss:.4} | hb[0]={:.4} | hw[0,0]={:.4} | b_cfc[0]={:.4}",
                heads_b[0], heads_w[0], b_cfc[0]
            );
        }
        prev_ts_ns = ts_ns;
        ts_ns += 20_000_000;
        prev_mid = 0.5 * (input.bid_px[0] + input.ask_px[0]);
    }
}

#[test]
fn debug_dump_grad_and_weight_movement() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        n_in: 32,
        n_hid: 32,
        seq_len: 1,
        lr: 3e-2,
        weight_decay: 0.0,
        seed: 0x4242,
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
    let stream = dev.cuda_stream().unwrap().clone();

    let snap_heads_b_init = download(&stream, &trainer.heads_b_d);
    let snap_heads_w_init = download(&stream, &trainer.heads_w_d);
    let snap_b_init = download(&stream, &trainer.b_d);
    println!("heads_b INITIAL: {:?}", snap_heads_b_init);
    println!("heads_w[0..5] INITIAL: {:?}", &snap_heads_w_init[0..5]);
    println!("b[0..5] INITIAL: {:?}", &snap_b_init[0..5]);

    let mut ts_ns = 1_000_000u64;
    let mut prev_ts_ns = 0u64;
    let mut prev_mid = 5500.0_f32;
    for step in 0..5 {
        trainer.reset_hidden_state().expect("reset");
        let (input, labels) = fixed_up_input(prev_mid, ts_ns, prev_ts_ns);
        let loss = trainer.step(&input, &labels).expect("step");

        let grad_heads_w = download(&stream, &trainer.heads_w_d); // weights after step
        let grad_heads_b = download(&stream, &trainer.heads_b_d);
        let b_post = download(&stream, &trainer.b_d);

        let dw0: f32 = grad_heads_w[0] - snap_heads_w_init[0];
        let db0: f32 = grad_heads_b[0] - snap_heads_b_init[0];
        let dbcfc0: f32 = b_post[0] - snap_b_init[0];

        println!(
            "step {step}: loss={loss:.4} | Δheads_w[0]={dw0:+.4e} | Δheads_b[0]={db0:+.4e} | Δb_cfc[0]={dbcfc0:+.4e}"
        );

        prev_ts_ns = ts_ns;
        ts_ns += 20_000_000;
        prev_mid = 0.5 * (input.bid_px[0] + input.ask_px[0]);
    }

    let snap_heads_b_final = download(&stream, &trainer.heads_b_d);
    let snap_heads_w_final = download(&stream, &trainer.heads_w_d);
    let snap_b_final = download(&stream, &trainer.b_d);

    // If the chain works on a constant-positive signal, heads_b[0] (the
    // bias for h=30 horizon) should have moved upward (positive grad
    // descent toward y=1). Floors to catch silent no-ops:
    let dheads_b_norm: f32 = snap_heads_b_final
        .iter()
        .zip(&snap_heads_b_init)
        .map(|(a, b)| (a - b).powi(2))
        .sum::<f32>()
        .sqrt();
    let dheads_w_norm: f32 = snap_heads_w_final
        .iter()
        .zip(&snap_heads_w_init)
        .map(|(a, b)| (a - b).powi(2))
        .sum::<f32>()
        .sqrt();
    let db_cfc_norm: f32 = snap_b_final
        .iter()
        .zip(&snap_b_init)
        .map(|(a, b)| (a - b).powi(2))
        .sum::<f32>()
        .sqrt();

    println!("Σ|Δheads_b|={dheads_b_norm:.4e}, Σ|Δheads_w|={dheads_w_norm:.4e}, Σ|Δb_cfc|={db_cfc_norm:.4e}");

    assert!(dheads_b_norm > 1e-4, "heads_b did not move — backward chain broken");
    assert!(dheads_w_norm > 1e-4, "heads_w did not move — backward chain broken");
    assert!(db_cfc_norm > 1e-6, "cfc_b did not move — cfc_step backward chain broken");
}