//! Output smoothness kernel — forward + gradient finite-diff validation.
//!
//! Per `feedback_no_cpu_test_fallbacks.md`: GPU oracle only. The analytic
//! gradient is verified against finite-difference of the kernel's own
//! forward pass (the kernel is the reference; we only check that its
//! analytic backward matches its own numerical derivative).
//!
//! Per `pearl_tests_must_prove_not_lock_observations.md`: assertions
//! verify INVARIANTS (loss ≥ 0 for non-zero diffs; loss == 0 when probs
//! are constant in k; finite-diff match) — not specific numeric outputs
//! that would lock the kernel to a single implementation.

use std::sync::Arc;

use anyhow::{Context, Result};
use approx::assert_relative_eq;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use ml_alpha::heads::N_HORIZONS;
use ml_core::device::MlDevice;

const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/output_smoothness.cubin"));

struct SmoothOut {
    loss: f32,
    raw_per_h: Vec<f32>,
    grad: Vec<f32>,
}

/// Launch the smoothness kernel against `probs [K*B*5]`, returning
/// (total λ-weighted loss, per-horizon raw mean-sq-diff, accumulated grad).
/// `grad_init` is the value `grad_probs` is pre-loaded with BEFORE the
/// kernel runs (since the kernel adds +=). Pass zeros to verify the
/// pure-smoothness gradient; non-zero values verify additivity.
fn run_smoothness_kernel(
    dev: &MlDevice,
    probs: &[f32],
    lambda_per_h: &[f32; N_HORIZONS],
    k: usize,
    b: usize,
    grad_init: &[f32],
) -> Result<SmoothOut> {
    let total = k * b * N_HORIZONS;
    assert_eq!(probs.len(), total);
    assert_eq!(grad_init.len(), total);

    let stream: &Arc<CudaStream> = dev.cuda_stream().context("smooth stream")?;
    let ctx = dev.cuda_context().context("smooth ctx")?;
    let module = ctx.load_cubin(CUBIN.to_vec()).context("load output_smoothness cubin")?;
    let func = module.load_function("output_smoothness_loss_and_grad").context("load fn")?;

    let probs_d   = upload(stream, probs)?;
    let lambda_d  = upload(stream, lambda_per_h.as_slice())?;
    let mut grad_d = upload(stream, grad_init)?;
    let mut loss_d = stream.alloc_zeros::<f32>(1).context("loss alloc")?;
    let mut raw_d  = stream.alloc_zeros::<f32>(N_HORIZONS).context("raw alloc")?;

    let k_i = k as i32;
    let b_i = b as i32;
    let cfg = LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = stream.launch_builder(&func);
    launch
        .arg(&probs_d).arg(&lambda_d)
        .arg(&k_i).arg(&b_i)
        .arg(&mut loss_d).arg(&mut raw_d).arg(&mut grad_d);
    unsafe { launch.launch(cfg).context("smoothness launch")?; }
    stream.synchronize().context("smoothness sync")?;

    Ok(SmoothOut {
        loss: download(stream, &loss_d)?[0],
        raw_per_h: download(stream, &raw_d)?,
        grad: download(stream, &grad_d)?,
    })
}

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    use ml_alpha::pinned_mem::MappedF32Buffer;
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("smooth upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream.alloc_zeros::<f32>(n).context("smooth upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream())
                .context("smooth upload DtoD")?;
        }
    }
    Ok(dst)
}

fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
    use ml_alpha::pinned_mem::MappedF32Buffer;
    let n = src.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("smooth download staging: {e}"))?;
    let nbytes = n * std::mem::size_of::<f32>();
    unsafe {
        let (src_ptr, _g) = src.device_ptr(stream);
        cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream())
            .context("smooth download DtoD")?;
    }
    stream.synchronize().context("smooth download sync")?;
    Ok(staging.read_all())
}

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

#[test]
fn smoothness_zero_when_probs_constant_in_k() {
    let dev = test_device();
    let k = 4usize;
    let b = 2usize;
    // probs[k, b, h] = 0.3 + 0.1*h + 0.05*b   (depends on b, h — NOT k)
    let mut probs = vec![0.0; k * b * N_HORIZONS];
    for kk in 0..k {
        for bb in 0..b {
            for h in 0..N_HORIZONS {
                probs[kk * b * N_HORIZONS + bb * N_HORIZONS + h] =
                    0.3 + 0.1 * (h as f32) + 0.05 * (bb as f32);
            }
        }
    }
    let lambda = [0.5_f32; N_HORIZONS];
    let grad_init = vec![0.0_f32; k * b * N_HORIZONS];
    let out = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &grad_init).unwrap();
    // Constant-in-k probs ⇒ raw_h[h] == 0 for all h ⇒ total loss == 0 and grad all zero.
    for (h, &raw) in out.raw_per_h.iter().enumerate() {
        assert!(raw.abs() < 1e-7, "raw_per_h[{h}] = {raw} (expected ~0 for k-constant probs)");
    }
    assert!(out.loss.abs() < 1e-7, "loss = {} (expected ~0)", out.loss);
    for (i, &g) in out.grad.iter().enumerate() {
        assert!(g.abs() < 1e-7, "grad[{i}] = {g} (expected ~0)");
    }
}

#[test]
fn smoothness_loss_positive_for_jittery_probs() {
    let dev = test_device();
    let k = 6usize;
    let b = 1usize;
    // probs alternate sharply in k for every (b, h).
    let mut probs = vec![0.0; k * b * N_HORIZONS];
    for kk in 0..k {
        for h in 0..N_HORIZONS {
            probs[kk * b * N_HORIZONS + h] = if kk % 2 == 0 { 0.2 } else { 0.8 };
        }
    }
    let lambda = [1.0_f32; N_HORIZONS];
    let grad_init = vec![0.0_f32; probs.len()];
    let out = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &grad_init).unwrap();
    // Each (p[k]-p[k-1])² = 0.36 for K-1=5 pairs per horizon ⇒ raw_h = 0.36 for all h.
    for (h, &raw) in out.raw_per_h.iter().enumerate() {
        assert!(raw > 0.3 && raw < 0.4, "raw_per_h[{h}] = {raw} (expected ≈ 0.36 for alternating 0.2/0.8)");
    }
    assert!(out.loss > 0.0, "loss should be positive for jittery probs, got {}", out.loss);
}

#[test]
fn smoothness_grad_is_additive_on_top_of_existing() {
    let dev = test_device();
    let k = 4usize;
    let b = 1usize;
    let probs: Vec<f32> = (0..k * b * N_HORIZONS).map(|i| 0.1 + 0.07 * (i as f32 % 7.0)).collect();
    let lambda = [0.3_f32; N_HORIZONS];

    // Run once with zero init, once with a known non-zero init.
    let zero_init = vec![0.0_f32; probs.len()];
    let out_zero = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &zero_init).unwrap();

    let nz_init: Vec<f32> = (0..probs.len()).map(|i| 0.5 + 0.01 * (i as f32)).collect();
    let out_nz = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &nz_init).unwrap();

    // out_nz.grad[i] should equal nz_init[i] + out_zero.grad[i].
    for i in 0..probs.len() {
        let expected = nz_init[i] + out_zero.grad[i];
        assert_relative_eq!(out_nz.grad[i], expected, epsilon = 5e-5, max_relative = 5e-4);
    }
}

#[test]
fn smoothness_grad_matches_finite_difference() {
    let dev = test_device();
    let k = 5usize;
    let b = 2usize;
    let mut probs = vec![0.0; k * b * N_HORIZONS];
    for i in 0..probs.len() {
        // Vary in k so derivatives are non-trivial.
        probs[i] = 0.2 + 0.05 * ((i % 7) as f32) + 0.03 * ((i / 7) as f32 % 5.0);
    }
    let lambda: [f32; N_HORIZONS] = [0.1, 1.0, 10.0];  // mimic the production ratio (3 horizons)
    let zero_init = vec![0.0_f32; probs.len()];
    let base = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &zero_init).unwrap();

    let eps = 1e-3_f32;
    // Probe a representative set of indices: corner (k=0), interior (k=2),
    // far corner (k=K-1), across two batches and horizons.
    let probe_indices: Vec<usize> = vec![
        // (k=0, b=0, h=0)
        0,
        // (k=0, b=1, h=2)
        0 * b * N_HORIZONS + 1 * N_HORIZONS + 2,
        // (k=2, b=0, h=1)
        2 * b * N_HORIZONS + 0 * N_HORIZONS + 1,
        // (k=2, b=1, h=4)
        2 * b * N_HORIZONS + 1 * N_HORIZONS + 4,
        // (k=K-1, b=0, h=3)
        (k - 1) * b * N_HORIZONS + 0 * N_HORIZONS + 3,
        // (k=K-1, b=1, h=0)
        (k - 1) * b * N_HORIZONS + 1 * N_HORIZONS + 0,
    ];
    for &i in &probe_indices {
        let mut up = probs.clone();
        up[i] += eps;
        let mut dn = probs.clone();
        dn[i] -= eps;
        let lu = run_smoothness_kernel(&dev, &up, &lambda, k, b, &zero_init).unwrap().loss;
        let ld = run_smoothness_kernel(&dev, &dn, &lambda, k, b, &zero_init).unwrap().loss;
        let fd = (lu - ld) / (2.0 * eps);
        assert_relative_eq!(
            base.grad[i],
            fd,
            epsilon = 5e-4,
            max_relative = 5e-2
        );
    }
}

#[test]
fn smoothness_zero_lambda_gives_zero_grad_and_loss() {
    // Default `smoothness_base_lambda = 0.0` ⇒ λ[h] = 0 ∀h ⇒ kernel
    // produces neither loss nor gradient. Verifies the opt-in default
    // is exactly behaviour-preserving.
    let dev = test_device();
    let k = 4usize;
    let b = 2usize;
    let probs: Vec<f32> = (0..k * b * N_HORIZONS).map(|i| 0.1 + 0.05 * (i as f32 % 9.0)).collect();
    let lambda = [0.0_f32; N_HORIZONS];
    let grad_init = vec![0.0_f32; probs.len()];
    let out = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &grad_init).unwrap();
    assert!(out.loss.abs() < 1e-7);
    for (i, &g) in out.grad.iter().enumerate() {
        assert!(g.abs() < 1e-7, "grad[{i}] = {g} (expected 0 for λ=0)");
    }
    // raw_per_h is still populated (telemetry — independent of λ).
    // We don't assert specific values; only that the loss is zero.
}
