// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

//! ndarray-style ergonomics: random init, scalar/value access, pretty-print.
//! Run: `cargo test -p rlx-tensor --features eval`.
#![cfg(feature = "eval")]

use rlx_tensor::Tensor;

#[test]
fn dims_and_numel() {
    let x = Tensor::zeros([2, 3, 4]);
    assert_eq!(x.dims(), vec![2, 3, 4]);
    assert_eq!(x.numel(), 24);
}

#[test]
fn item_reads_scalar() {
    let s = (&Tensor::from_vec(vec![3.0, 4.0], [2]) * &Tensor::from_vec(vec![3.0, 4.0], [2]))
        .sum([0], false); // 9 + 16 = 25
    assert_eq!(s.item(), 25.0);
}

#[test]
fn rng_is_reproducible_and_distinct() {
    // Same seed -> identical; different seed -> different.
    let a = Tensor::randn([100], 42).to_vec();
    let b = Tensor::randn([100], 42).to_vec();
    let c = Tensor::randn([100], 7).to_vec();
    assert_eq!(a, b, "same seed must reproduce");
    assert_ne!(a, c, "different seed must differ");
}

#[test]
fn randn_stats_are_sane() {
    let n = 10_000;
    let v = Tensor::randn([n], 123).to_vec();
    let mean: f32 = v.iter().sum::<f32>() / n as f32;
    let var: f32 = v.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / n as f32;
    assert!(mean.abs() < 0.05, "mean ~0, got {mean}");
    assert!((var - 1.0).abs() < 0.1, "var ~1, got {var}");
}

#[test]
fn rand_range_bounds() {
    let v = Tensor::rand_range([5000], 9, -2.0, 3.0).to_vec();
    assert!(
        v.iter().all(|&x| (-2.0..3.0).contains(&x)),
        "out of [lo,hi)"
    );
    let max = v.iter().cloned().fold(f32::MIN, f32::max);
    let min = v.iter().cloned().fold(f32::MAX, f32::min);
    assert!(max > 2.0 && min < -1.0, "should span most of the range");
}

#[test]
fn randn_init_trains_mlp_like() {
    // Random init is usable for a tiny fit: y = sum(w*x), learn w toward x sign.
    // Here just sanity: randn weights produce finite forward output.
    let w = Tensor::randn_std([3], 5, 0.0, 0.5);
    let x = Tensor::from_vec(vec![1.0, 2.0, 3.0], [3]);
    let out = (&w * &x).sum([0], false).to_vec();
    assert!(out[0].is_finite());
}

#[cfg(feature = "eval-metal")]
#[test]
fn rng_is_backend_agnostic() {
    use rlx_tensor::{Device, fastest_device, is_available};
    assert!(is_available(Device::Metal));
    assert_eq!(fastest_device(), Device::Metal);
    // Host-generated random constants must produce identical values on every
    // backend (they're baked as Op::Constant; the backend only reads bytes).
    let normal = Tensor::randn([256], 2024);
    assert_eq!(
        normal.to_vec_on(Device::Cpu),
        normal.to_vec_on(Device::Metal),
        "randn must match across CPU and Metal"
    );
    let uniform = Tensor::rand_range([256], 99, -1.0, 1.0);
    assert_eq!(
        uniform.to_vec_on(Device::Cpu),
        uniform.to_vec_on(Device::Metal),
        "rand_range must match across CPU and Metal"
    );
}

#[test]
fn pretty_print() {
    let x = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], [2, 2]);
    let s = x.show();
    assert!(s.contains("Tensor[2, 2]"), "header: {s}");
    assert!(s.contains("[[1, 2], [3, 4]]"), "nested values: {s}");
    // Display delegates to show().
    assert_eq!(format!("{x}"), s);
}
