ferrotorch 0.6.2

//! Correctness validation: verify ferrotorch produces the same results
//! as PyTorch for core operations.
//!
//! Each `eprintln!("PASS: …")` follows a real `assert!` / `assert_close` /
//! `assert_eq!` and is informational — using stderr means the diagnostic is
//! visible in test output without `--nocapture`, while the actual assertion
//! is what gates the test.
//!
//! Run with: cargo test --test validate_vs_pytorch --release

use std::time::Instant;

use ferrotorch_core::*;
use ferrotorch_nn::*;
use ferrotorch_optim::*;

fn assert_close(a: f32, b: f32, tol: f32, msg: &str) {
    assert!(
        (a - b).abs() < tol,
        "{msg}: {a} vs {b} (diff = {})",
        (a - b).abs()
    );
}

// =====================================================================
// CORRECTNESS: Softmax sums to 1
// =====================================================================
#[test]
fn test_softmax_sums_to_one() {
    let x = randn::<f32>(&[4, 10]).unwrap();
    let sm = grad_fns::activation::softmax(&x).unwrap();
    let sm_data = sm.data().unwrap();

    for row in 0..4 {
        let row_sum: f32 = (0..10).map(|c| sm_data[row * 10 + c]).sum();
        assert_close(row_sum, 1.0, 1e-5, &format!("softmax row {row} sum"));
    }
    eprintln!("PASS: softmax sums to 1");
}

// =====================================================================
// CORRECTNESS: LayerNorm zero mean, unit variance
// =====================================================================
#[test]
fn test_layernorm_statistics() {
    let ln = LayerNorm::new(vec![256], 1e-5, false).unwrap();
    // Input with non-zero mean and non-unit variance
    let x = (&randn::<f32>(&[4, 32, 256]).unwrap() * &scalar::<f32>(5.0).unwrap()).unwrap();
    let x = (&x + &scalar::<f32>(3.0).unwrap()).unwrap();
    let y = ln.forward(&x).unwrap();
    let y_data = y.data().unwrap();

    let batch_seq = 4 * 32;
    let d = 256;
    for i in 0..batch_seq {
        let slice = &y_data[i * d..(i + 1) * d];
        let mean: f32 = slice.iter().sum::<f32>() / d as f32;
        let var: f32 = slice.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / d as f32;
        assert!(mean.abs() < 1e-4, "layernorm mean too large: {mean}");
        assert!(
            (var - 1.0).abs() < 1e-3,
            "layernorm var too far from 1: {var}"
        );
    }
    eprintln!("PASS: layernorm produces zero mean, unit variance");
}

// =====================================================================
// CORRECTNESS: Autograd numerical gradient check (x^2)
// =====================================================================
#[test]
fn test_autograd_x_squared() {
    let x = from_vec(vec![1.0f32, 2.0, 3.0], &[3])
        .unwrap()
        .requires_grad_(true);
    let x2 = (&x * &x).unwrap();
    let loss = x2.sum_all().unwrap();
    loss.backward().unwrap();

    let grad = x.grad().unwrap().unwrap();
    let grad_data = grad.data().unwrap();
    // d/dx(x^2) = 2x
    assert_close(grad_data[0], 2.0, 1e-4, "grad[0]");
    assert_close(grad_data[1], 4.0, 1e-4, "grad[1]");
    assert_close(grad_data[2], 6.0, 1e-4, "grad[2]");
    eprintln!("PASS: autograd d/dx(x^2) = 2x");
}

// =====================================================================
// CORRECTNESS: CrossEntropy produces gradients
// =====================================================================
#[test]
fn test_crossentropy_gradient() {
    let logits = randn::<f32>(&[4, 10]).unwrap().requires_grad_(true);
    let targets = from_vec(vec![3.0, 7.0, 1.0, 5.0], &[4]).unwrap();
    let ce = CrossEntropyLoss::new(Reduction::Mean, 0.0);
    let loss = ce.forward(&logits, &targets).unwrap();
    loss.backward().unwrap();

    let grad = logits.grad().unwrap().unwrap();
    let grad_data = grad.data().unwrap();
    let grad_sum: f32 = grad_data.iter().map(|g| g.abs()).sum();
    assert!(grad_sum > 0.0, "CrossEntropy gradient should be non-zero");
    eprintln!("PASS: CrossEntropy produces non-zero gradients (sum = {grad_sum:.4})");
}

// =====================================================================
// CORRECTNESS: Dropout masks at correct rate
// =====================================================================
#[test]
fn test_dropout_rate() {
    let x = ones::<f32>(&[10000]).unwrap();
    let dropout = Dropout::new(0.3).unwrap();
    // Training mode
    let y = dropout.forward(&x).unwrap();
    let y_data = y.data().unwrap();
    let zero_count = y_data.iter().filter(|&&v| v == 0.0).count();
    let zero_frac = zero_count as f32 / 10000.0;
    assert!(
        zero_frac > 0.20 && zero_frac < 0.40,
        "dropout rate should be ~0.3, got {zero_frac}"
    );
    eprintln!("PASS: dropout zero fraction = {zero_frac:.4} (expected ~0.3)");
}

// =====================================================================
// CORRECTNESS: Embedding lookup matches weight rows
// =====================================================================
#[test]
fn test_embedding_lookup() {
    let emb = Embedding::new(100, 16, None).unwrap();
    let ids = from_vec(vec![5.0f32, 10.0, 50.0], &[3]).unwrap();
    let out = emb.forward(&ids).unwrap();
    let out_data = out.data().unwrap();

    let weight_data = emb.parameters()[0].tensor().data().unwrap();
    // Row 5 should match out[0]
    for j in 0..16 {
        assert_close(
            out_data[j],
            weight_data[5 * 16 + j],
            1e-6,
            &format!("emb[5][{j}]"),
        );
    }
    // Row 50 should match out[2]
    for j in 0..16 {
        assert_close(
            out_data[2 * 16 + j],
            weight_data[50 * 16 + j],
            1e-6,
            &format!("emb[50][{j}]"),
        );
    }
    eprintln!("PASS: embedding lookup matches weight rows");
}

// =====================================================================
// CORRECTNESS: Causal attention mask
// =====================================================================
#[test]
fn test_causal_attention_mask() {
    let seq = 8;
    let d = 16;
    let q = randn::<f32>(&[1, seq, d]).unwrap();
    let k = randn::<f32>(&[1, seq, d]).unwrap();

    // Q @ K^T
    let k_t = permute_t(&k, &[0, 2, 1]).unwrap();
    let scores = grad_fns::linalg::bmm_differentiable(&q, &k_t).unwrap();

    // Causal mask
    let mask_data: Vec<f32> = (0..seq * seq)
        .map(|i| if (i % seq) <= (i / seq) { 0.0 } else { -1e9 })
        .collect();
    let mask = from_vec(mask_data, &[1, seq, seq]).unwrap();
    let masked_scores = (&scores + &mask).unwrap();

    let attn = grad_fns::activation::softmax(&masked_scores).unwrap();
    let attn_data = attn.data().unwrap();

    // Upper triangle should be ~0 (future positions masked)
    let mut upper_sum = 0.0f32;
    for i in 0..seq {
        for j in (i + 1)..seq {
            upper_sum += attn_data[i * seq + j].abs();
        }
    }
    assert!(
        upper_sum < 1e-5,
        "causal mask: upper triangle sum should be ~0, got {upper_sum}"
    );
    eprintln!("PASS: causal attention mask zeros future positions (upper sum = {upper_sum:.2e})");
}

// =====================================================================
// CORRECTNESS: Conv2d output shape
// =====================================================================
#[test]
fn test_conv2d_output_shape() {
    let conv = Conv2d::new(3, 16, (3, 3), (2, 2), (1, 1), true).unwrap();
    let x = rand::<f32>(&[2, 3, 32, 32]).unwrap();
    let y = conv.forward(&x).unwrap();
    assert_eq!(y.shape(), &[2, 16, 16, 16], "conv2d output shape mismatch");
    eprintln!("PASS: Conv2d [2,3,32,32] stride=2 pad=1 -> [2,16,16,16]");
}

// =====================================================================
// TASK 1: MLP training convergence
// =====================================================================
#[test]
fn test_mlp_training_converges() {
    let model = Sequential::new(vec![
        Box::new(Linear::new(20, 64, true).unwrap()),
        Box::new(ReLU::default()),
        Box::new(Linear::new(64, 5, true).unwrap()),
    ]);

    let mut adam_cfg = AdamConfig::default();
    adam_cfg.lr = 1e-3;
    let mut optimizer = Adam::new(model.parameters().into_iter().cloned().collect(), adam_cfg);

    let ce = CrossEntropyLoss::new(Reduction::Mean, 0.0);
    let mut losses = Vec::new();

    // Fixed data so loss reliably decreases (random data each step is noisy)
    let x_fixed = randn::<f32>(&[16, 20]).unwrap();
    let target_fixed = from_vec((0..16).map(|i| (i % 5) as f32).collect(), &[16]).unwrap();

    for step in 0..20 {
        let x = &x_fixed;
        let target = &target_fixed;

        optimizer.zero_grad().unwrap();
        let output = model.forward(x).unwrap();
        let loss = ce.forward(&output, target).unwrap();
        let loss_val = loss.data_vec().unwrap()[0];
        losses.push(loss_val);

        loss.backward().unwrap();

        let model_params: Vec<&Parameter<f32>> = model.parameters();
        let opt_params = optimizer.param_groups()[0].params();
        for (mp, op) in model_params.iter().zip(opt_params.iter()) {
            if let Some(g) = mp.grad().unwrap() {
                op.set_grad(Some(g)).unwrap();
            }
        }
        optimizer.step().unwrap();

        if step == 0 || step == 19 {
            eprintln!("  Step {step}: loss = {loss_val:.4}");
        }
    }

    assert!(
        losses.last().unwrap() < losses.first().unwrap(),
        "MLP loss should decrease: first={}, last={}",
        losses.first().unwrap(),
        losses.last().unwrap()
    );
    eprintln!(
        "PASS: MLP training loss decreases ({:.4} -> {:.4})",
        losses[0], losses[19]
    );
}

// =====================================================================
// TASK 2: Autoencoder reconstruction
// =====================================================================
#[test]
fn test_autoencoder_reconstruction() {
    let encoder = Sequential::new(vec![
        Box::new(Linear::new(784, 256, true).unwrap()),
        Box::new(ReLU::default()),
        Box::new(Linear::new(256, 32, true).unwrap()),
    ]);
    let decoder = Sequential::new(vec![
        Box::new(Linear::new(32, 256, true).unwrap()),
        Box::new(ReLU::default()),
        Box::new(Linear::new(256, 784, true).unwrap()),
    ]);

    let mut all_params: Vec<Parameter<f32>> = Vec::new();
    all_params.extend(encoder.parameters().into_iter().cloned());
    all_params.extend(decoder.parameters().into_iter().cloned());

    let mut adam_cfg = AdamConfig::default();
    adam_cfg.lr = 1e-3;
    let mut optimizer = Adam::new(all_params, adam_cfg);

    let mut losses = Vec::new();

    for step in 0..20 {
        let x = randn::<f32>(&[16, 784]).unwrap();
        optimizer.zero_grad().unwrap();

        let encoded = encoder.forward(&x).unwrap();
        let decoded = decoder.forward(&encoded).unwrap();

        // MSE loss
        let diff = (&decoded - &x).unwrap();
        let sq = (&diff * &diff).unwrap();
        let loss = sq.mean_all().unwrap();
        let loss_val = loss.data_vec().unwrap()[0];
        losses.push(loss_val);

        loss.backward().unwrap();

        // Sync grads
        let enc_params: Vec<&Parameter<f32>> = encoder.parameters();
        let dec_params: Vec<&Parameter<f32>> = decoder.parameters();
        let all_model_params: Vec<&Parameter<f32>> =
            enc_params.into_iter().chain(dec_params).collect();
        let opt_params = optimizer.param_groups()[0].params();
        for (mp, op) in all_model_params.iter().zip(opt_params.iter()) {
            if let Some(g) = mp.grad().unwrap() {
                op.set_grad(Some(g)).unwrap();
            }
        }
        optimizer.step().unwrap();

        if step == 0 || step == 19 {
            eprintln!("  Step {step}: recon_loss = {loss_val:.4}");
        }
    }

    assert!(
        losses.last().unwrap() < losses.first().unwrap(),
        "Autoencoder loss should decrease: first={}, last={}",
        losses.first().unwrap(),
        losses.last().unwrap()
    );
    eprintln!(
        "PASS: Autoencoder reconstruction loss decreases ({:.4} -> {:.4})",
        losses[0], losses[19]
    );
}

// =====================================================================
// TASK 3: Transformer block training
// =====================================================================
#[test]
fn test_transformer_block_training() {
    let d_model = 64;
    let n_heads = 4;
    let head_dim = d_model / n_heads;
    let batch = 2;
    let seq = 8;

    let qkv_proj = Linear::new(d_model, 3 * d_model, false).unwrap();
    let out_proj = Linear::new(d_model, d_model, false).unwrap();
    let ln = LayerNorm::new(vec![d_model], 1e-5, true).unwrap();

    let mut all_params: Vec<Parameter<f32>> = Vec::new();
    all_params.extend(qkv_proj.parameters().into_iter().cloned());
    all_params.extend(out_proj.parameters().into_iter().cloned());
    all_params.extend(ln.parameters().into_iter().cloned());

    let mut adamw_cfg = AdamWConfig::default();
    adamw_cfg.lr = 3e-3;
    let mut optimizer = AdamW::new(all_params, adamw_cfg);
    let mut losses = Vec::new();

    // Use fixed input/target so loss can decrease (random data each step is harder)
    let x_fixed = randn::<f32>(&[batch, seq, d_model]).unwrap();
    let target_fixed = randn::<f32>(&[batch, seq, d_model]).unwrap();

    for step in 0..30 {
        let x = &x_fixed;
        let target = &target_fixed;

        optimizer.zero_grad().unwrap();

        let normed = ln.forward(x).unwrap();
        let qkv = qkv_proj.forward(&normed).unwrap();
        let qkv_chunks = chunk_t(&qkv, 3, 2).unwrap();

        let q = qkv_chunks[0]
            .view(&[batch as i64 * n_heads as i64, seq as i64, head_dim as i64])
            .unwrap();
        let k = qkv_chunks[1]
            .view(&[batch as i64 * n_heads as i64, seq as i64, head_dim as i64])
            .unwrap();
        let v = qkv_chunks[2]
            .view(&[batch as i64 * n_heads as i64, seq as i64, head_dim as i64])
            .unwrap();

        let k_t = permute_t(&k, &[0, 2, 1]).unwrap();
        let scores = grad_fns::linalg::bmm_differentiable(&q, &k_t).unwrap();
        let scale = scalar::<f32>((head_dim as f32).sqrt()).unwrap();
        let scores = grad_fns::arithmetic::div(&scores, &scale).unwrap();
        let attn = grad_fns::activation::softmax(&scores).unwrap();
        let attn_out = grad_fns::linalg::bmm_differentiable(&attn, &v).unwrap();
        let attn_out = attn_out
            .view(&[batch as i64, seq as i64, d_model as i64])
            .unwrap();
        let output = out_proj.forward(&attn_out).unwrap();

        // Residual
        let output = (x + &output).unwrap();

        // MSE loss vs target
        let diff = (&output - target).unwrap();
        let loss = (&diff * &diff).unwrap().mean_all().unwrap();
        let loss_val = loss.data_vec().unwrap()[0];
        losses.push(loss_val);

        loss.backward().unwrap();

        let model_params: Vec<&Parameter<f32>> = qkv_proj
            .parameters()
            .into_iter()
            .chain(out_proj.parameters())
            .chain(ln.parameters())
            .collect();
        let opt_params = optimizer.param_groups()[0].params();
        for (mp, op) in model_params.iter().zip(opt_params.iter()) {
            if let Some(g) = mp.grad().unwrap() {
                op.set_grad(Some(g)).unwrap();
            }
        }
        optimizer.step().unwrap();

        if step == 0 || step == 29 {
            eprintln!("  Step {step}: loss = {loss_val:.4}");
        }
    }

    assert!(
        losses.last().unwrap() < losses.first().unwrap(),
        "Transformer loss should decrease: first={}, last={}",
        losses.first().unwrap(),
        losses.last().unwrap()
    );
    eprintln!(
        "PASS: Transformer block training loss decreases ({:.4} -> {:.4})",
        losses[0], losses[29]
    );
}

// =====================================================================
// SPEED: Timed comparison benchmarks
//
// This test asserts perf-ratio thresholds against a PyTorch CPU reference so
// that gross regressions (e.g. accidental debug-build, dropped SIMD, or
// quadratic-in-rank rewrites of a hot kernel) surface as a hard test
// failure rather than as a quiet log line. Thresholds are deliberately
// generous: the goal is to catch order-of-magnitude regressions, not to
// tune to within 20% of the reference on any machine. `#[ignore]`-gated
// because timing on a busy CI runner is noisy and the test must be run on
// a quiet host (`--release --ignored`).
// =====================================================================
#[test]
#[ignore = "perf-ratio benchmark; noisy on shared CI runners — run explicitly with `--release -- --ignored`"]
fn test_speed_comparison() {
    eprintln!("\n{}", "=".repeat(60));
    eprintln!("SPEED COMPARISON vs PyTorch");
    eprintln!("{}", "=".repeat(60));

    // PyTorch CPU reference numbers, in microseconds per op, captured by
    // `scripts/pytorch_validate.py` on the project reference machine.
    let pytorch_ref: &[(&str, f64)] = &[
        ("add_1M", 91.7),
        ("mul_1M", 41.9),
        ("relu_1M", 19.9),
        ("sigmoid_1M", 163.5),
        ("matmul_64", 5.2),
        ("matmul_256", 145.4),
        ("matmul_1024", 2853.0),
    ];

    // Regression guard: ferrotorch must stay within this multiple of the
    // PyTorch reference. Picked from the observed worst-case ratio on a
    // developer laptop (release mode): elementwise `mul_1M` lands around
    // ~19x of PyTorch today because PyTorch dispatches a hand-tuned AVX
    // kernel while ferrotorch goes through the generic ferray loop.
    // 25x leaves headroom for run-to-run jitter on noisy hosts and still
    // catches an order-of-magnitude regression (e.g. accidentally
    // shipping a debug build, dropping SIMD, or recomputing a transpose
    // in a hot loop — any of which would push the ratio well past 50x).
    const MAX_RATIO_VS_PYTORCH: f64 = 25.0;

    let a = rand::<f32>(&[1000, 1000]).unwrap();
    let b = rand::<f32>(&[1000, 1000]).unwrap();

    let run = |_name: &str, iters: usize, f: &dyn Fn()| -> f64 {
        for _ in 0..5 {
            f();
        }
        let start = Instant::now();
        for _ in 0..iters {
            f();
        }

        start.elapsed().as_secs_f64() / iters as f64 * 1e6
    };

    let ft_add = run("add_1M", 100, &|| {
        let _ = (&a + &b).unwrap();
    });
    let ft_mul = run("mul_1M", 100, &|| {
        let _ = (&a * &b).unwrap();
    });
    let ft_relu = run("relu_1M", 100, &|| {
        let _ = a.relu().unwrap();
    });
    let ft_sig = run("sigmoid_1M", 100, &|| {
        let _ = a.sigmoid().unwrap();
    });

    let a64 = rand::<f32>(&[64, 64]).unwrap();
    let b64 = rand::<f32>(&[64, 64]).unwrap();
    let ft_mm64 = run("matmul_64", 100, &|| {
        let _ = a64.matmul(&b64).unwrap();
    });

    let a256 = rand::<f32>(&[256, 256]).unwrap();
    let b256 = rand::<f32>(&[256, 256]).unwrap();
    let ft_mm256 = run("matmul_256", 100, &|| {
        let _ = a256.matmul(&b256).unwrap();
    });

    let a1024 = rand::<f32>(&[1024, 1024]).unwrap();
    let b1024 = rand::<f32>(&[1024, 1024]).unwrap();
    let ft_mm1024 = run("matmul_1024", 20, &|| {
        let _ = a1024.matmul(&b1024).unwrap();
    });

    let results: &[(&str, f64)] = &[
        ("add_1M", ft_add),
        ("mul_1M", ft_mul),
        ("relu_1M", ft_relu),
        ("sigmoid_1M", ft_sig),
        ("matmul_64", ft_mm64),
        ("matmul_256", ft_mm256),
        ("matmul_1024", ft_mm1024),
    ];

    eprintln!(
        "\n{:<20} {:>12} {:>14} {:>8}",
        "Operation", "PyTorch(us)", "ferrotorch(us)", "Ratio"
    );
    eprintln!("{}", "-".repeat(58));

    let mut failures: Vec<String> = Vec::new();
    for (name, ft_us) in results {
        let pt_us = pytorch_ref
            .iter()
            .find(|(n, _)| n == name)
            .map(|(_, v)| *v)
            .unwrap_or_else(|| panic!("missing PyTorch reference for {name}"));
        let ratio = ft_us / pt_us;
        eprintln!(
            "{:<20} {:>12.1} {:>14.1} {:>7.1}x",
            name, pt_us, ft_us, ratio
        );
        if ratio > MAX_RATIO_VS_PYTORCH {
            failures.push(format!(
                "{name}: ferrotorch={ft_us:.1}us vs PyTorch={pt_us:.1}us \
                 (ratio={ratio:.1}x, threshold={MAX_RATIO_VS_PYTORCH:.1}x)"
            ));
        }
    }

    // Final hard gate. All ops must clear the threshold; we collect every
    // offender before failing so a single run surfaces all regressions.
    assert!(
        failures.is_empty(),
        "perf-ratio regression vs PyTorch reference:\n  - {}",
        failures.join("\n  - ")
    );
}