//! BCE forward + backward validation.
//!
//! Per `feedback_no_cpu_test_fallbacks.md`: no CPU oracle. The gradient
//! is validated against finite-difference computed via GPU forward calls
//! at perturbed inputs (the kernel itself is the reference; we only
//! verify that the analytic grad matches its own numerical derivative).

use approx::assert_relative_eq;
use ml_alpha::trainer::loss::{bce_multi_horizon_loss_and_grad_gpu, BceInput};
use ml_core::device::MlDevice;

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

fn make_input(probs: Vec<f32>, labels: Vec<f32>) -> BceInput {
    BceInput { probs, labels, n_horizons: 5, n_pos: 4 }
}

#[test]
fn bce_loss_is_positive_and_finite() {
    let dev = test_device();
    let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
    let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
    let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
    assert!(out.loss.is_finite());
    assert!(out.loss > 0.0);
    assert_eq!(out.n_valid, 20);
}

#[test]
fn bce_grad_zero_when_probs_match_labels() {
    let dev = test_device();
    // probs ≈ labels (within clamp window): grad ≈ 0 wherever |p-y| ≈ 0
    let mut probs = vec![0.5_f32; 20];
    let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
    // Move probs very close to labels (but within (1e-6, 1-1e-6) clamp window).
    for i in 0..20 {
        probs[i] = if labels[i] > 0.5 { 1.0 - 1e-5 } else { 1e-5 };
    }
    let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
    // Near-perfect predictions → near-zero loss
    assert!(out.loss < 1e-3, "loss should be small for matching probs/labels, got {}", out.loss);
}

#[test]
fn bce_grad_matches_finite_difference() {
    let dev = test_device();
    let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
    let labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
    let base = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs.clone(), labels.clone())).unwrap();

    let eps = 1e-3_f32;
    for i in [0usize, 5, 10, 15, 19] {
        let mut up = probs.clone();
        up[i] += eps;
        let mut dn = probs.clone();
        dn[i] -= eps;
        let lu = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(up, labels.clone())).unwrap().loss;
        let ld = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(dn, labels.clone())).unwrap().loss;
        let fd = (lu - ld) / (2.0 * eps);
        assert_relative_eq!(
            base.grad_probs[i],
            fd,
            epsilon = 5e-3,
            max_relative = 5e-2
        );
    }
}

#[test]
fn bce_nan_labels_mask_grad_and_loss() {
    let dev = test_device();
    let mut labels: Vec<f32> = (0..20).map(|i| (i % 2) as f32).collect();
    // Mask out indices 5, 10, 15
    labels[5] = f32::NAN;
    labels[10] = f32::NAN;
    labels[15] = f32::NAN;
    let probs: Vec<f32> = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect();
    let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap();
    assert_eq!(out.n_valid, 17, "should drop 3 NaN-masked entries");
    for &i in &[5usize, 10, 15] {
        assert_eq!(out.grad_probs[i], 0.0, "masked grad should be 0 at index {i}, got {}", out.grad_probs[i]);
    }
}