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

//! Cross-validation tests against `llama.cpp` reference dequantization.
//!
//! These tests load pre-computed fixtures generated by
//! `tests/fixtures/gguf_reference/generate_gguf.py` and compare anamnesis
//! `GGUF` block-quant dequantization output against the `gguf` Python
//! package's `dequantize` function — the official reference implementation
//! from `ggml-org` that mirrors `ggml-quants.c`'s `dequantize_row_*`
//! functions.
//!
//! Each fixture is a 65 536-element slice (2 048 blocks for legacy quants,
//! 256 super-blocks for K-quants) extracted from a real model tensor.
//!
//! Coverage: **22 of 22** production kernels from 4 real models + 3 synthetic
//! fixtures (`Q4_0`–`Q8_0`, `Q2_K`–`Q6_K`, `IQ4_NL`, `IQ4_XS`, `IQ2_XXS`,
//! `IQ2_XS`, `IQ2_S`, `IQ3_XXS`, `IQ3_S`, `IQ1_S`, `IQ1_M`, `TQ1_0`, `TQ2_0`,
//! `MXFP4`). `TQ1_0` / `TQ2_0` / `MXFP4` are synthetic-fixture cross-validated
//! — only ~15 BitNet-derivative GGUFs ship the `TQ*` types on `HuggingFace`,
//! and mainstream `MXFP4` only ships inside the 11 GB `gpt-oss-20b` upload;
//! Python `gguf.quants.quantize()` implements all three encode sides, so a
//! deterministic random tensor is the practical fixture source. `Q8_1` and
//! `Q8_K` are not shipped by any real model — they are internal `llama.cpp`
//! activation quant types, already covered by unit tests.

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

use std::time::Instant;

use anamnesis::{dequantize_gguf_to_bf16, GgufType};

// ---------------------------------------------------------------------------
// Fixture parsing
// ---------------------------------------------------------------------------

/// Binary fixture layout (all little-endian):
///
/// - 4 bytes: `ggml_type` discriminant (`u32`)
/// - 4 bytes: `n_elements` (`u32`)
/// - 4 bytes: raw data byte count (`u32`)
/// - 4 bytes: golden `BF16` byte count (`u32`)
/// - raw quantized block data
/// - expected `BF16` output
struct GgufFixture {
    n_elements: usize,
    raw_data: Vec<u8>,
    expected_bf16: Vec<u8>,
}

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

/// Maps a `ggml_type` discriminant to a [`GgufType`].
///
/// Covers all 22 dequantizable block types — Phase 4.5 step 6 closed
/// the GGUF coverage gap, so every block-quantised `GgufType` variant
/// is exercised here.
fn gguf_type_from_disc(disc: u32) -> GgufType {
    match disc {
        2 => GgufType::Q4_0,
        3 => GgufType::Q4_1,
        6 => GgufType::Q5_0,
        7 => GgufType::Q5_1,
        8 => GgufType::Q8_0,
        9 => GgufType::Q8_1,
        10 => GgufType::Q2_K,
        11 => GgufType::Q3_K,
        12 => GgufType::Q4_K,
        13 => GgufType::Q5_K,
        14 => GgufType::Q6_K,
        15 => GgufType::Q8_K,
        20 => GgufType::IQ4_NL,
        23 => GgufType::IQ4_XS,
        16 => GgufType::IQ2_XXS,
        17 => GgufType::IQ2_XS,
        22 => GgufType::IQ2_S,
        18 => GgufType::IQ3_XXS,
        21 => GgufType::IQ3_S,
        19 => GgufType::IQ1_S,
        29 => GgufType::IQ1_M,
        34 => GgufType::TQ1_0,
        35 => GgufType::TQ2_0,
        39 => GgufType::MXFP4,
        other => panic!("unknown ggml_type discriminant: {other}"),
    }
}

fn parse_gguf_fixture(data: &[u8], expected_dtype: GgufType) -> GgufFixture {
    let disc = read_u32_le(data, 0);
    let n_elements = read_u32_le(data, 4) as usize;
    let raw_data_len = read_u32_le(data, 8) as usize;
    let golden_len = read_u32_le(data, 12) as usize;

    let actual_dtype = gguf_type_from_disc(disc);
    assert_eq!(
        actual_dtype, expected_dtype,
        "fixture dtype mismatch: expected {expected_dtype:?}, got {actual_dtype:?} (disc={disc})"
    );

    let header_size = 16;
    let raw_start = header_size;
    let golden_start = raw_start + raw_data_len;

    GgufFixture {
        n_elements,
        raw_data: data[raw_start..raw_start + raw_data_len].to_vec(),
        expected_bf16: data[golden_start..golden_start + golden_len].to_vec(),
    }
}

// ---------------------------------------------------------------------------
// BF16 comparison
// ---------------------------------------------------------------------------

/// Compare two `BF16` byte slices, allowing up to `max_ulp_diff` ULP
/// (unit in the last place) difference per element.
///
/// Returns the number of mismatched elements and the maximum ULP diff found.
fn compare_bf16(actual: &[u8], expected: &[u8], max_ulp_diff: u16) -> (usize, u16) {
    assert_eq!(actual.len(), expected.len(), "output length mismatch");
    let mut mismatches = 0;
    let mut max_diff: u16 = 0;

    for (i, (a_pair, e_pair)) in actual
        .chunks_exact(2)
        .zip(expected.chunks_exact(2))
        .enumerate()
    {
        let a_bits = u16::from_le_bytes([a_pair[0], a_pair[1]]);
        let e_bits = u16::from_le_bytes([e_pair[0], e_pair[1]]);

        // Handle NaN: both NaN is a match.
        let a_is_nan = (a_bits & 0x7F80 == 0x7F80) && (a_bits & 0x007F != 0);
        let e_is_nan = (e_bits & 0x7F80 == 0x7F80) && (e_bits & 0x007F != 0);
        if a_is_nan && e_is_nan {
            continue;
        }
        if a_is_nan != e_is_nan {
            mismatches += 1;
            continue;
        }

        let diff = a_bits.abs_diff(e_bits);
        if diff > max_ulp_diff {
            mismatches += 1;
            if i < 5 {
                eprintln!(
                    "  element {i}: actual=0x{a_bits:04X}, expected=0x{e_bits:04X}, diff={diff} ULP"
                );
            }
        }
        if diff > max_diff {
            max_diff = diff;
        }
    }
    (mismatches, max_diff)
}

// ---------------------------------------------------------------------------
// Unified test runner
// ---------------------------------------------------------------------------

fn run_cross_validation(name: &str, data: &[u8], dtype: GgufType, max_ulp: u16) {
    let fixture = parse_gguf_fixture(data, dtype);
    let total = fixture.n_elements;

    // Dequantize with anamnesis and measure time.
    let start = Instant::now();
    let actual = dequantize_gguf_to_bf16(&fixture.raw_data, dtype, fixture.n_elements)
        .expect("dequantization failed");
    let elapsed = start.elapsed();

    assert_eq!(
        actual.len(),
        fixture.expected_bf16.len(),
        "{name}: output length mismatch"
    );

    let (mismatches, max_diff) = compare_bf16(&actual, &fixture.expected_bf16, max_ulp);
    eprintln!(
        "{name}: {total} elements, {mismatches} mismatches, \
         max ULP diff = {max_diff}, anamnesis = {:.1} \u{b5}s",
        elapsed.as_secs_f64() * 1e6
    );
    assert_eq!(
        mismatches, 0,
        "{name}: {mismatches}/{total} elements differ by more than {max_ulp} ULP"
    );
}

// ---------------------------------------------------------------------------
// Legacy quants (block_size=32) — 6 kernels (incl. IQ4_NL)
// ---------------------------------------------------------------------------

#[test]
fn cross_validate_smollm2_q4_0() {
    run_cross_validation(
        "SmolLM2-135M Q4_0 (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q4_0.bin"),
        GgufType::Q4_0,
        0,
    );
}

#[test]
fn cross_validate_smollm2_q4_1() {
    run_cross_validation(
        "SmolLM2-135M Q4_1 (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q4_1.bin"),
        GgufType::Q4_1,
        0,
    );
}

#[test]
fn cross_validate_tinyllama_q5_0() {
    run_cross_validation(
        "TinyLlama-1.1B Q5_0 (TheBloke)",
        include_bytes!("fixtures/gguf_reference/tinyllama_q5_0.bin"),
        GgufType::Q5_0,
        0,
    );
}

#[test]
fn cross_validate_smollm2_q5_1() {
    run_cross_validation(
        "SmolLM2-135M Q5_1 (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q5_1.bin"),
        GgufType::Q5_1,
        0,
    );
}

#[test]
fn cross_validate_smollm2_q8_0() {
    run_cross_validation(
        "SmolLM2-135M Q8_0 (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q8_0.bin"),
        GgufType::Q8_0,
        0,
    );
}

#[test]
fn cross_validate_smollm2_iq4_nl() {
    run_cross_validation(
        "SmolLM2-135M IQ4_NL (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_iq4_nl.bin"),
        GgufType::IQ4_NL,
        0,
    );
}

// ---------------------------------------------------------------------------
// K-quants (block_size=256) — 6 kernels (incl. IQ4_XS)
// ---------------------------------------------------------------------------

#[test]
fn cross_validate_tinyllama_q2_k() {
    run_cross_validation(
        "TinyLlama-1.1B Q2_K (TheBloke)",
        include_bytes!("fixtures/gguf_reference/tinyllama_q2_k.bin"),
        GgufType::Q2_K,
        0,
    );
}

#[test]
fn cross_validate_smollm2_q3_k() {
    run_cross_validation(
        "SmolLM2-135M Q3_K (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q3_k.bin"),
        GgufType::Q3_K,
        0,
    );
}

#[test]
fn cross_validate_smollm2_q4_k() {
    run_cross_validation(
        "SmolLM2-135M Q4_K (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q4_k.bin"),
        GgufType::Q4_K,
        0,
    );
}

#[test]
fn cross_validate_smollm2_q5_k() {
    run_cross_validation(
        "SmolLM2-135M Q5_K (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q5_k.bin"),
        GgufType::Q5_K,
        0,
    );
}

#[test]
fn cross_validate_smollm2_q6_k() {
    run_cross_validation(
        "SmolLM2-135M Q6_K (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_q6_k.bin"),
        GgufType::Q6_K,
        0,
    );
}

#[test]
fn cross_validate_smollm2_iq4_xs() {
    run_cross_validation(
        "SmolLM2-135M IQ4_XS (bartowski)",
        include_bytes!("fixtures/gguf_reference/smollm2_iq4_xs.bin"),
        GgufType::IQ4_XS,
        0,
    );
}

// ---------------------------------------------------------------------------
// IQ2 variants (256-element super-blocks) — 3 kernels
// ---------------------------------------------------------------------------
//
// Small models ship `IQ2_M` as the only IQ2 mix, which internally uses
// `IQ2_S` for most critical layers and no `IQ2_XXS`/`IQ2_XS` at all. The
// smallest repo shipping pure `IQ2_XXS.gguf` and `IQ2_XS.gguf` files is
// `bartowski/Mistral-7B-Instruct-v0.3-GGUF`; `IQ2_S` fixtures come from
// `bartowski/Qwen2.5-0.5B-Instruct-GGUF`'s `IQ2_M` mix (21 `IQ2_S` tensors
// inside).

#[test]
fn cross_validate_mistral_7b_iq2_xxs() {
    run_cross_validation(
        "Mistral-7B-v0.3 IQ2_XXS (bartowski)",
        include_bytes!("fixtures/gguf_reference/mistral_7b_iq2_xxs.bin"),
        GgufType::IQ2_XXS,
        0,
    );
}

#[test]
fn cross_validate_mistral_7b_iq2_xs() {
    run_cross_validation(
        "Mistral-7B-v0.3 IQ2_XS (bartowski)",
        include_bytes!("fixtures/gguf_reference/mistral_7b_iq2_xs.bin"),
        GgufType::IQ2_XS,
        0,
    );
}

#[test]
fn cross_validate_qwen25_iq2_s() {
    run_cross_validation(
        "Qwen2.5-0.5B IQ2_S (bartowski)",
        include_bytes!("fixtures/gguf_reference/qwen25_iq2_s.bin"),
        GgufType::IQ2_S,
        0,
    );
}

// ---------------------------------------------------------------------------
// IQ3 variants (256-element super-blocks) — 2 kernels
// ---------------------------------------------------------------------------
//
// Both fixtures extracted from a single `bartowski/Mistral-7B-Instruct-v0.3-
// GGUF / Mistral-7B-Instruct-v0.3-IQ3_XXS.gguf` download — that file's
// primary type is `IQ3_XXS` (96 tensors) but it also ships 33 `IQ3_S`
// tensors as the high-precision secondary, so one 2.64 GB fetch covers
// both step 3 fixtures.

#[test]
fn cross_validate_mistral_7b_iq3_xxs() {
    run_cross_validation(
        "Mistral-7B-v0.3 IQ3_XXS (bartowski)",
        include_bytes!("fixtures/gguf_reference/mistral_7b_iq3_xxs.bin"),
        GgufType::IQ3_XXS,
        0,
    );
}

#[test]
fn cross_validate_mistral_7b_iq3_s() {
    run_cross_validation(
        "Mistral-7B-v0.3 IQ3_S (bartowski)",
        include_bytes!("fixtures/gguf_reference/mistral_7b_iq3_s.bin"),
        GgufType::IQ3_S,
        0,
    );
}

// ---------------------------------------------------------------------------
// IQ1 variants (256-element super-blocks) — 2 kernels
// ---------------------------------------------------------------------------
//
// Smallest IQ-family quants. Both source files come from
// `bartowski/Mistral-7B-Instruct-v0.3-GGUF` but, unlike IQ3, the IQ1
// variants don't share a single file: `Mistral-7B-Instruct-v0.3-IQ1_S.gguf`
// has 156 IQ1_S tensors but no IQ1_M; `...-IQ1_M.gguf` has 156 IQ1_M
// tensors but no IQ1_S. Two downloads, two fixtures.

#[test]
fn cross_validate_mistral_7b_iq1_s() {
    run_cross_validation(
        "Mistral-7B-v0.3 IQ1_S (bartowski)",
        include_bytes!("fixtures/gguf_reference/mistral_7b_iq1_s.bin"),
        GgufType::IQ1_S,
        0,
    );
}

#[test]
fn cross_validate_mistral_7b_iq1_m() {
    run_cross_validation(
        "Mistral-7B-v0.3 IQ1_M (bartowski)",
        include_bytes!("fixtures/gguf_reference/mistral_7b_iq1_m.bin"),
        GgufType::IQ1_M,
        0,
    );
}

// ---------------------------------------------------------------------------
// TQ variants (256-element super-blocks) — 2 kernels, synthetic fixtures
// ---------------------------------------------------------------------------
//
// Ternary quantizations invented for BitNet-style 1.58-bit models. Only ~15
// community uploads exist on HuggingFace; mainstream model providers don't
// ship TQ variants. Python `gguf.quants.quantize()` implements both, so the
// fixtures come from a deterministic synthetic random tensor (seed=42,
// scale=0.1) — see `generate_gguf.py::generate_synthetic_fixture`. Same
// `gguf.dequantize()` reference path as real-model fixtures, just a
// different input source.

#[test]
fn cross_validate_synthetic_tq1_0() {
    run_cross_validation(
        "Synthetic TQ1_0 (seed=42)",
        include_bytes!("fixtures/gguf_reference/synthetic_tq1_0.bin"),
        GgufType::TQ1_0,
        0,
    );
}

#[test]
fn cross_validate_synthetic_tq2_0() {
    run_cross_validation(
        "Synthetic TQ2_0 (seed=42)",
        include_bytes!("fixtures/gguf_reference/synthetic_tq2_0.bin"),
        GgufType::TQ2_0,
        0,
    );
}

// ---------------------------------------------------------------------------
// MXFP4 (32-element block) — synthetic fixture, completes Phase 4.5
// ---------------------------------------------------------------------------
//
// Microscaling FP4 (OCP MX standard from late 2023, added to ggml in 2024).
// Mainstream MXFP4 GGUFs only ship inside the 11 GB `gpt-oss-20b` upload —
// too large to justify when Python `gguf.quants.quantize()` supports MXFP4
// and a deterministic synthetic random tensor (seed=42, scale=0.1) gives
// bit-exact cross-validation. Reuses the synthetic-fixture path established
// for TQ1_0 / TQ2_0 in step 5.

#[test]
fn cross_validate_synthetic_mxfp4() {
    run_cross_validation(
        "Synthetic MXFP4 (seed=42)",
        include_bytes!("fixtures/gguf_reference/synthetic_mxfp4.bin"),
        GgufType::MXFP4,
        0,
    );
}
