//! multi_horizon_heads GPU invariants.
//!
//! Per `feedback_no_cpu_test_fallbacks.md`: validation via analytically-
//! known synthetic inputs + property assertions.

use approx::assert_relative_eq;
use ml_alpha::heads::{multi_horizon_heads_gpu, HeadsWeights, HIDDEN_DIM, N_HORIZONS};
use ml_core::device::MlDevice;

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

#[test]
fn probs_are_in_unit_interval() {
    let dev = test_device();
    let w = HeadsWeights {
        w: vec![0.1; N_HORIZONS * HIDDEN_DIM],
        b: vec![0.0; N_HORIZONS],
    };
    let h: Vec<f32> = (0..HIDDEN_DIM).map(|i| 0.01 * i as f32).collect();
    let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
    for k in 0..N_HORIZONS {
        assert!((0.0..=1.0).contains(&probs[k]), "head {k} prob {} not in [0,1]", probs[k]);
    }
}

#[test]
fn zero_weights_and_bias_give_half() {
    // logit = 0 -> sigmoid(0) = 0.5
    let dev = test_device();
    let w = HeadsWeights {
        w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
        b: vec![0.0; N_HORIZONS],
    };
    let h: Vec<f32> = (0..HIDDEN_DIM).map(|i| i as f32).collect();
    let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
    for k in 0..N_HORIZONS {
        assert_relative_eq!(probs[k], 0.5_f32, epsilon = 1e-6);
    }
}

#[test]
fn large_positive_bias_saturates_high() {
    let dev = test_device();
    let w = HeadsWeights {
        w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
        b: vec![20.0; N_HORIZONS],
    };
    let h = vec![0.0; HIDDEN_DIM];
    let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
    for k in 0..N_HORIZONS {
        assert!(probs[k] > 0.999, "head {k} should saturate near 1 with bias=20, got {}", probs[k]);
    }
}

#[test]
fn large_negative_bias_saturates_low() {
    let dev = test_device();
    let w = HeadsWeights {
        w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
        b: vec![-20.0; N_HORIZONS],
    };
    let h = vec![0.0; HIDDEN_DIM];
    let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
    for k in 0..N_HORIZONS {
        assert!(probs[k] < 0.001, "head {k} should saturate near 0 with bias=-20, got {}", probs[k]);
    }
}

#[test]
fn per_head_independence() {
    // Per `feedback_use_consts_not_literals_for_structural_dims`:
    // address heads by N_HORIZONS-relative offsets, not hardcoded
    // indices that silently drift when N_HORIZONS changes.
    // Set first bias = +5, middle biases = 0, last bias = -5; with
    // zero weights and h=0 the head output is sigmoid(bias):
    //   sigmoid(+5) ≈ 0.993, sigmoid(0) = 0.5, sigmoid(-5) ≈ 0.0067.
    assert!(N_HORIZONS >= 3, "test requires at least 3 horizons");
    let dev = test_device();
    let mut b = vec![0.0_f32; N_HORIZONS];
    b[0] = 5.0;
    b[N_HORIZONS - 1] = -5.0;
    let w = HeadsWeights {
        w: vec![0.0; N_HORIZONS * HIDDEN_DIM],
        b,
    };
    let h = vec![0.0; HIDDEN_DIM];
    let probs = multi_horizon_heads_gpu(&dev, &w, &h).expect("gpu");
    assert!(
        probs[0] > 0.99 && probs[0] < 1.0,
        "probs[0]=sigmoid(+5) should be > 0.99; got {}",
        probs[0]
    );
    for k in 1..N_HORIZONS - 1 {
        assert_relative_eq!(probs[k], 0.5, epsilon = 1e-6);
    }
    let last = N_HORIZONS - 1;
    assert!(
        probs[last] > 0.0 && probs[last] < 0.01,
        "probs[{last}]=sigmoid(-5) should be in (0, 0.01); got {}",
        probs[last]
    );
}
