// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Bulk validation: run the INT8 model over a few hundred MNIST test
//! images and assert the accuracy stays close to the FP32 baseline.
//!
//! The test set is generated by `tools/train_mnist.py` (along with the
//! weights) and lives at `tests/data/test_set.bin`. Format: for each
//! image, 784 bytes i8 NHWC followed by 1 byte label.

use rlx_cortexm::model::{INPUT_LEN, SCRATCH_LEN, infer};
use rlx_cortexm::model_weights as w;

const REC_LEN: usize = INPUT_LEN + 1;

/// FP32 baseline is ~97%. Min accuracy floor is set per weight
/// bit-width: i8 / i4 are within 1-2 pp of fp32, but ternary PTQ
/// without QAT collapses to ~44% (the network can't represent
/// enough of its weights as `{-1, 0, 1}`). The infrastructure is
/// still bit-exact with the host numpy reference at i2 — the test
/// just lowers the floor to "is the kernel running" instead of
/// "is accuracy preserved" because preservation isn't possible
/// without retraining. See README.md for the full discussion.
const fn min_acc_for_bits(bits: u8) -> f64 {
    match bits {
        2 => 0.20, // PTQ ternary; sanity floor only
        _ => 0.92,
    }
}

fn data_path() -> std::path::PathBuf {
    let mut p = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    p.push("tests/data/test_set.bin");
    p
}

#[test]
fn int8_accuracy_close_to_fp32() {
    let path = data_path();
    let blob = match std::fs::read(&path) {
        Ok(b) => b,
        Err(_) => {
            eprintln!(
                "skipping: {} not found — run tools/train_mnist.py",
                path.display()
            );
            return;
        }
    };
    assert_eq!(blob.len() % REC_LEN, 0, "test_set.bin malformed");
    let n = blob.len() / REC_LEN;

    let mut a = vec![0i8; SCRATCH_LEN];
    let mut b = vec![0i8; SCRATCH_LEN];
    let mut correct = 0usize;

    for i in 0..n {
        let rec = &blob[i * REC_LEN..(i + 1) * REC_LEN];
        let img: &[i8] = unsafe {
            std::slice::from_raw_parts(rec[..INPUT_LEN].as_ptr() as *const i8, INPUT_LEN)
        };
        let label = rec[INPUT_LEN];
        let pred = infer(img, &mut a, &mut b);
        if pred as u8 == label {
            correct += 1;
        }
    }

    let acc = correct as f64 / n as f64;
    let floor = min_acc_for_bits(w::WEIGHT_BITS);
    println!(
        "INT8 accuracy on {n} images (weight_bits={}): {correct}/{n} = {:.4}",
        w::WEIGHT_BITS,
        acc
    );
    assert!(
        acc >= floor,
        "INT8 accuracy {acc:.4} below floor {floor:.4} (weight_bits={})",
        w::WEIGHT_BITS
    );
}
