//! CUDA Graph A capture + replay.
//!
//! Validates that the captured graph produces identical output to the
//! sequential dispatch path for the SAME input (scalars are frozen at
//! capture time per cudarc 0.19 semantics — see the trunk struct doc).

use approx::assert_relative_eq;
use ml_alpha::cfc::snap_features::Mbp10RawInput;
use ml_alpha::cfc::{CfcConfig, CfcTrunk};
use ml_alpha::heads::{N_HORIZONS, PROJ_DIM};
use ml_core::device::MlDevice;

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

fn synthetic_input(dt_ns: u64) -> Mbp10RawInput {
    let mut bid_px = [0.0f32; 10];
    let mut bid_sz = [0.0f32; 10];
    let mut ask_px = [0.0f32; 10];
    let mut ask_sz = [0.0f32; 10];
    for i in 0..10 {
        bid_px[i] = 5500.00 - 0.25 * i as f32;
        ask_px[i] = 5500.25 + 0.25 * i as f32;
        bid_sz[i] = 12.0 + i as f32;
        ask_sz[i] = 10.0 + i as f32 * 1.5;
    }
    Mbp10RawInput {
        bid_px, bid_sz, ask_px, ask_sz,
        prev_mid: 5499.875,
        trade_signed_vol: 4.0,
        trade_count: 7,
        ts_ns: dt_ns,
        prev_ts_ns: 0,
        regime: [0.0; 6],
    }
}

#[test]
fn graph_a_replay_matches_sequential() {
    let dev = test_device();
    let input = synthetic_input(20_000_000);

    // Trunk A: sequential dispatch.
    let mut trunk_seq = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xCAFE).expect("seq init");
    let (probs_seq, proj_seq) = trunk_seq.forward_snapshot(&input).expect("seq forward");

    // Trunk B: same seed, capture + replay.
    let mut trunk_g = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xCAFE).expect("graph init");
    trunk_g.capture_graph_a(&input).expect("capture");
    let (probs_g, proj_g) = trunk_g.perception_forward_captured().expect("replay");

    for k in 0..N_HORIZONS {
        assert_relative_eq!(probs_seq[k], probs_g[k], epsilon = 1e-5, max_relative = 1e-5);
    }
    for j in 0..PROJ_DIM {
        assert_relative_eq!(proj_seq[j], proj_g[j], epsilon = 1e-4, max_relative = 1e-4);
    }
}

#[test]
fn graph_a_replay_is_deterministic() {
    let dev = test_device();
    let input = synthetic_input(15_000_000);
    let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0x5EED).expect("init");
    trunk.capture_graph_a(&input).expect("capture");

    let (p1, _) = trunk.perception_forward_captured().expect("replay 1");
    let (p2, _) = trunk.perception_forward_captured().expect("replay 2");
    let (p3, _) = trunk.perception_forward_captured().expect("replay 3");

    for k in 0..N_HORIZONS {
        assert_eq!(p1[k], p2[k], "replay {k} differs between calls 1 and 2");
        assert_eq!(p2[k], p3[k], "replay {k} differs between calls 2 and 3");
    }
}

#[test]
fn graph_a_replay_outputs_finite() {
    let dev = test_device();
    let input = synthetic_input(40_000_000);
    let mut trunk = CfcTrunk::new_random(&dev, &CfcConfig::default(), 0xBEEF).expect("init");
    trunk.capture_graph_a(&input).expect("capture");
    let (probs, proj) = trunk.perception_forward_captured().expect("replay");
    for &p in &probs {
        assert!(p.is_finite() && (0.0..=1.0).contains(&p));
    }
    for &v in &proj {
        assert!(v.is_finite());
    }
}
