// SPDX-License-Identifier: MIT OR Apache-2.0

//! Cross-validation tests for `NPZ` parsing against `NumPy` reference.
//!
//! Validates `parse_npz` output against fixtures generated from the Gemma Scope
//! 2B SAE weights (`params.npz`, 5 `F32` arrays). The Python script
//! `tests/fixtures/npz_reference/generate_npz.py` creates the fixtures.

#![cfg(feature = "npz")]
#![allow(
    clippy::panic,
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::indexing_slicing,
    clippy::as_conversions,
    clippy::cast_possible_truncation,
    clippy::similar_names,
    clippy::wildcard_enum_match_arm,
    clippy::float_cmp
)]

use std::time::Instant;

use anamnesis::{parse_npz, NpzDtype};

// ---------------------------------------------------------------------------
// Reference fixture parsing
// ---------------------------------------------------------------------------

/// A single array's expected metadata and spot-check data from the binary
/// fixture generated by `generate_npz.py`.
struct RefArray {
    name: String,
    shape: Vec<u64>,
    dtype_id: u32,
    data: Vec<u8>,
}

/// Parse the binary reference fixture into a list of expected arrays.
///
/// Format (all LE): see `generate_npz.py` docstring.
fn parse_reference(data: &[u8]) -> Vec<RefArray> {
    let mut offset = 0;

    let num_arrays = read_u32(data, &mut offset);
    let mut arrays = Vec::with_capacity(num_arrays as usize);

    for _ in 0..num_arrays {
        let name_len = read_u32(data, &mut offset) as usize;
        let name = std::str::from_utf8(&data[offset..offset + name_len])
            .unwrap()
            .to_owned();
        offset += name_len;

        let num_dims = read_u32(data, &mut offset) as usize;
        let mut shape = Vec::with_capacity(num_dims);
        for _ in 0..num_dims {
            shape.push(read_u64(data, &mut offset));
        }

        let dtype_id = read_u32(data, &mut offset);

        let data_len = read_u32(data, &mut offset) as usize;
        let arr_data = data[offset..offset + data_len].to_vec();
        offset += data_len;

        arrays.push(RefArray {
            name,
            shape,
            dtype_id,
            data: arr_data,
        });
    }

    arrays
}

fn read_u32(data: &[u8], offset: &mut usize) -> u32 {
    let val = u32::from_le_bytes(data[*offset..*offset + 4].try_into().unwrap());
    *offset += 4;
    val
}

fn read_u64(data: &[u8], offset: &mut usize) -> u64 {
    let val = u64::from_le_bytes(data[*offset..*offset + 8].try_into().unwrap());
    *offset += 8;
    val
}

/// Map `NpzDtype` to the fixture's dtype ID (must match `generate_npz.py`).
fn dtype_to_id(dtype: NpzDtype) -> u32 {
    match dtype {
        NpzDtype::Bool => 0,
        NpzDtype::U8 => 1,
        NpzDtype::I8 => 2,
        NpzDtype::U16 => 3,
        NpzDtype::I16 => 4,
        NpzDtype::U32 => 5,
        NpzDtype::I32 => 6,
        NpzDtype::U64 => 7,
        NpzDtype::I64 => 8,
        NpzDtype::F16 => 9,
        NpzDtype::BF16 => 10,
        NpzDtype::F32 => 11,
        NpzDtype::F64 => 12,
        _ => panic!("unknown NpzDtype variant"),
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

/// Validate that `parse_npz` correctly reads a small `NPZ` generated from
/// the Gemma Scope 2B SAE weights. Checks array count, names, shapes,
/// dtypes, and byte-exact data match against the `NumPy` reference.
#[test]
fn cross_validate_gemma_scope_sae() {
    let fixture_path = concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/tests/fixtures/npz_reference/gemma_scope_small.npz"
    );
    let reference_data = include_bytes!("fixtures/npz_reference/gemma_scope_reference.bin");

    // Parse reference fixture
    let expected = parse_reference(reference_data);
    assert_eq!(expected.len(), 5, "expected 5 arrays in reference");

    // Parse NPZ with anamnesis
    let start = Instant::now();
    let tensors = parse_npz(fixture_path).expect("parse_npz failed");
    let elapsed = start.elapsed();

    eprintln!(
        "Gemma Scope SAE (small): {} arrays, parsed in {:.1} µs",
        tensors.len(),
        elapsed.as_secs_f64() * 1e6
    );

    assert_eq!(
        tensors.len(),
        expected.len(),
        "array count mismatch: got {}, expected {}",
        tensors.len(),
        expected.len()
    );

    // Verify each array against reference
    for ref_arr in &expected {
        let tensor = tensors
            .get(&ref_arr.name)
            .unwrap_or_else(|| panic!("array '{}' not found in parse_npz output", ref_arr.name));

        // Name (redundant field matches key)
        assert_eq!(tensor.name, ref_arr.name, "name mismatch");

        // Shape: the small NPZ has flattened 1D arrays [8] (truncated from the
        // original shapes stored in ref_arr.shape). Verify element count matches.
        let expected_elements = 8usize;
        let actual_elements: usize = tensor.shape.iter().product();
        assert_eq!(
            actual_elements, expected_elements,
            "element count mismatch for '{}': shape {:?} = {actual_elements}, expected {expected_elements}",
            ref_arr.name, tensor.shape
        );

        // Log original shape from reference for diagnostics
        eprintln!(
            "  original shape: {:?} (truncated to [{expected_elements}])",
            ref_arr.shape
        );

        // Dtype
        let actual_dtype_id = dtype_to_id(tensor.dtype);
        assert_eq!(
            actual_dtype_id, ref_arr.dtype_id,
            "dtype mismatch for '{}': got {actual_dtype_id}, expected {}",
            ref_arr.name, ref_arr.dtype_id
        );

        // Data: byte-exact match
        assert_eq!(
            tensor.data.len(),
            ref_arr.data.len(),
            "data length mismatch for '{}': got {}, expected {}",
            ref_arr.name,
            tensor.data.len(),
            ref_arr.data.len()
        );
        assert_eq!(
            tensor.data,
            ref_arr.data,
            "data mismatch for '{}': first differing byte at position {}",
            ref_arr.name,
            tensor
                .data
                .iter()
                .zip(ref_arr.data.iter())
                .position(|(a, b)| a != b)
                .unwrap_or(0)
        );

        eprintln!(
            "  {}: shape={:?}, dtype={}, {} bytes OK",
            ref_arr.name,
            tensor.shape,
            tensor.dtype,
            tensor.data.len()
        );
    }
}
