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

//! Cross-validation tests for `BitsAndBytes` dequantization against `PyTorch` reference.
//!
//! These tests load pre-computed fixtures generated by
//! `tests/fixtures/bnb_reference/generate_bnb.py` and compare anamnesis
//! `BnB` dequantization output against the `bitsandbytes` formula.

#![cfg(feature = "bnb")]
#![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
)]

use std::time::Instant;

use anamnesis::remember::bnb::dequantize_bnb4_double_quant_to_bf16;
use anamnesis::{dequantize_bnb4_to_bf16, dequantize_bnb_int8_to_bf16};

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

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)
}

/// Parsed `NF4`/`FP4` fixture (`format_id` = 0 or 2).
struct Bnb4Fixture {
    format_id: u32,
    total_elements: usize,
    block_size: usize,
    weight_data: Vec<u8>,
    absmax_data: Vec<u8>,
    quant_map_data: Vec<u8>,
    nested_absmax_data: Vec<u8>,
    nested_quant_map_data: Vec<u8>,
    expected_bf16: Vec<u8>,
}

/// Parsed `INT8` fixture (`format_id` = 1).
struct BnbInt8Fixture {
    out_features: usize,
    in_features: usize,
    weight_data: Vec<u8>,
    scb_data: Vec<u8>,
    expected_bf16: Vec<u8>,
}

fn parse_bnb4_fixture(data: &[u8]) -> Bnb4Fixture {
    let format_id = read_u32_le(data, 0);
    let total_elements = read_u32_le(data, 4) as usize;
    let block_size = read_u32_le(data, 8) as usize;
    let weight_len = read_u32_le(data, 12) as usize;
    let absmax_len = read_u32_le(data, 16) as usize;
    let quant_map_len = read_u32_le(data, 20) as usize;
    let nested_absmax_len = read_u32_le(data, 24) as usize;
    let nested_quant_map_len = read_u32_le(data, 28) as usize;
    let expected_len = read_u32_le(data, 32) as usize;

    let header_size = 36;
    let mut offset = header_size;

    let weight_data = data[offset..offset + weight_len].to_vec();
    offset += weight_len;
    let absmax_data = data[offset..offset + absmax_len].to_vec();
    offset += absmax_len;
    let quant_map_data = data[offset..offset + quant_map_len].to_vec();
    offset += quant_map_len;
    let nested_absmax_data = data[offset..offset + nested_absmax_len].to_vec();
    offset += nested_absmax_len;
    let nested_quant_map_data = data[offset..offset + nested_quant_map_len].to_vec();
    offset += nested_quant_map_len;
    let expected_bf16 = data[offset..offset + expected_len].to_vec();

    Bnb4Fixture {
        format_id,
        total_elements,
        block_size,
        weight_data,
        absmax_data,
        quant_map_data,
        nested_absmax_data,
        nested_quant_map_data,
        expected_bf16,
    }
}

fn parse_int8_fixture(data: &[u8]) -> BnbInt8Fixture {
    let _format_id = read_u32_le(data, 0); // = 1
    let out_features = read_u32_le(data, 4) as usize;
    let in_features = read_u32_le(data, 8) as usize;
    let weight_len = read_u32_le(data, 12) as usize;
    let scb_len = read_u32_le(data, 16) as usize;
    let expected_len = read_u32_le(data, 20) as usize;

    let header_size = 24;
    let mut offset = header_size;

    let weight_data = data[offset..offset + weight_len].to_vec();
    offset += weight_len;
    let scb_data = data[offset..offset + scb_len].to_vec();
    offset += scb_len;
    let expected_bf16 = data[offset..offset + expected_len].to_vec();

    BnbInt8Fixture {
        out_features,
        in_features,
        weight_data,
        scb_data,
        expected_bf16,
    }
}

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

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)
}

// ---------------------------------------------------------------------------
// Test runners
// ---------------------------------------------------------------------------

fn run_bnb4_cross_validation(name: &str, data: &[u8], max_ulp: u16) {
    let fixture = parse_bnb4_fixture(data);

    let format_label = match fixture.format_id {
        0 => "NF4/FP4",
        2 => "NF4/FP4 double-quant",
        other => panic!("unknown format_id: {other}"),
    };

    eprintln!(
        "{name}: {format_label}, block_size={}, {} elements",
        fixture.block_size, fixture.total_elements,
    );

    let start = Instant::now();
    let actual = if fixture.format_id == 2 {
        // Infer nested_block_size: absmax_count / nested_absmax_count
        let absmax_count = fixture.absmax_data.len();
        let nested_absmax_count = fixture.nested_absmax_data.len() / 4;
        let nested_block_size = if nested_absmax_count > 0 {
            absmax_count.div_ceil(nested_absmax_count)
        } else {
            256
        };
        dequantize_bnb4_double_quant_to_bf16(
            &fixture.weight_data,
            &fixture.absmax_data,
            &fixture.quant_map_data,
            &fixture.nested_absmax_data,
            &fixture.nested_quant_map_data,
            fixture.total_elements,
            fixture.block_size,
            nested_block_size,
        )
        .expect("BnB4 double-quant dequantization failed")
    } else {
        dequantize_bnb4_to_bf16(
            &fixture.weight_data,
            &fixture.absmax_data,
            &fixture.quant_map_data,
            fixture.total_elements,
            fixture.block_size,
        )
        .expect("BnB4 dequantization failed")
    };
    let elapsed = start.elapsed();

    let (mismatches, max_diff) = compare_bf16(&actual, &fixture.expected_bf16, max_ulp);
    eprintln!(
        "  {mismatches} mismatches, max ULP diff = {max_diff}, anamnesis = {:.1} µs",
        elapsed.as_secs_f64() * 1e6,
    );
    assert_eq!(
        mismatches, 0,
        "{name}: {mismatches} BF16 mismatches (max ULP diff = {max_diff})"
    );
}

fn run_int8_cross_validation(name: &str, data: &[u8], max_ulp: u16) {
    let fixture = parse_int8_fixture(data);

    eprintln!(
        "{name}: INT8, {}×{} = {} elements",
        fixture.out_features,
        fixture.in_features,
        fixture.out_features * fixture.in_features,
    );

    let start = Instant::now();
    let actual = dequantize_bnb_int8_to_bf16(
        &fixture.weight_data,
        &fixture.scb_data,
        fixture.out_features,
        fixture.in_features,
    )
    .expect("BnB INT8 dequantization failed");
    let elapsed = start.elapsed();

    let (mismatches, max_diff) = compare_bf16(&actual, &fixture.expected_bf16, max_ulp);
    eprintln!(
        "  {mismatches} mismatches, max ULP diff = {max_diff}, anamnesis = {:.1} µs",
        elapsed.as_secs_f64() * 1e6,
    );
    assert_eq!(
        mismatches, 0,
        "{name}: {mismatches} BF16 mismatches (max ULP diff = {max_diff})"
    );
}

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

#[test]
fn cross_validate_llama_1b_nf4() {
    let data = include_bytes!("fixtures/bnb_reference/llama_1b_nf4.bin");
    run_bnb4_cross_validation("Llama-3.2-1B NF4", data, 0);
}

#[test]
fn cross_validate_llama_1b_nf4_double_quant() {
    let data = include_bytes!("fixtures/bnb_reference/llama_1b_nf4_double_quant.bin");
    run_bnb4_cross_validation("Llama-3.2-1B NF4 double-quant", data, 0);
}

#[test]
fn cross_validate_llama_1b_fp4() {
    let data = include_bytes!("fixtures/bnb_reference/llama_1b_fp4.bin");
    run_bnb4_cross_validation("Llama-3.2-1B FP4", data, 0);
}

#[test]
fn cross_validate_llama_1b_int8() {
    let data = include_bytes!("fixtures/bnb_reference/llama_1b_int8.bin");
    run_int8_cross_validation("Llama-3.2-1B INT8", data, 0);
}
