//! Minimal repro: does candle's `xs.qmatmul(qt)` produce the same result as
//! `xs.matmul(qt.dequantize().t())` for various Q8_0 / Q4_K weights?
//!
//! If they diverge on a fresh random matrix → confirmed candle bug.
//! If they match on random but diverge for our checkpoint → numerical edge case.
//!
//! Run on Metal AND CPU separately:
//!   cargo run -p rtx-csm --release --features metal --example qmatmul_repro -- metal
//!   cargo run -p rtx-csm --release --features metal --example qmatmul_repro -- cpu

use anyhow::Result;
use candle_core::quantized::{GgmlDType, QMatMul, QTensor};
use candle_core::{DType, Device, Module, Tensor};

fn rms(a: &[f32], b: &[f32]) -> f32 {
    let n = a.len() as f32;
    let mut sum = 0.0f32;
    for (x, y) in a.iter().zip(b) {
        sum += (x - y).powi(2);
    }
    (sum / n).sqrt()
}

fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b)
        .map(|(x, y)| (x - y).abs())
        .fold(0.0f32, f32::max)
}

fn run_case(
    name: &str,
    weight: &Tensor,
    xs: &Tensor,
    dtype: GgmlDType,
    device: &Device,
) -> Result<()> {
    println!("\n=== {name} (qtype={dtype:?}) ===");
    println!(
        "weight shape={:?} dtype={:?}, xs shape={:?} dtype={:?}",
        weight.shape(),
        weight.dtype(),
        xs.shape(),
        xs.dtype()
    );

    // Quantize twice (QTensor isn't Clone): one for qmm path, one for deq.
    unsafe {
        std::env::remove_var("CANDLE_DEQUANTIZE_ALL");
        std::env::remove_var("CANDLE_DEQUANTIZE_ALL_F16");
    }
    let qt_a = QTensor::quantize(weight, dtype)?;
    let qt_b = QTensor::quantize(weight, dtype)?;
    println!("qt shape={:?} dtype={:?}", qt_a.shape(), qt_a.dtype());

    // Path A: QMatMul forward via raw QTensor variant — the suspect kernel.
    let qmm = QMatMul::from_qtensor(qt_a)?;
    let path_a = qmm.forward(xs)?;

    // Path B: dequantize manually, then xs @ deq.t() — the known-good path.
    let deq = qt_b.dequantize(device)?.to_dtype(xs.dtype())?;
    let path_b = xs.matmul(&deq.t()?)?;

    // Path C: matmul against the original (unquantized) weight — ground truth
    // mod the quantization error itself. Distance from C tells us how much
    // each path errs vs the float baseline.
    let path_c = xs.matmul(&weight.t()?)?;

    let a_vals = path_a
        .flatten_all()?
        .to_dtype(DType::F32)?
        .to_vec1::<f32>()?;
    let b_vals = path_b
        .flatten_all()?
        .to_dtype(DType::F32)?
        .to_vec1::<f32>()?;
    let c_vals = path_c
        .flatten_all()?
        .to_dtype(DType::F32)?
        .to_vec1::<f32>()?;

    let a_vs_b_rms = rms(&a_vals, &b_vals);
    let a_vs_b_max = max_abs_diff(&a_vals, &b_vals);
    let b_vs_c_rms = rms(&b_vals, &c_vals); // pure quant error
    let a_vs_c_rms = rms(&a_vals, &c_vals); // qmatmul error vs float baseline

    println!(
        "path A (xs.qmatmul) vs path B (xs @ deq.t):  rms={a_vs_b_rms:.6}  max={a_vs_b_max:.6}"
    );
    println!(
        "path B (dequant) vs path C (float baseline): rms={b_vs_c_rms:.6}  (pure quant noise)"
    );
    println!("path A (qmatmul) vs path C (float baseline): rms={a_vs_c_rms:.6}");
    if a_vs_b_rms > 10.0 * b_vs_c_rms {
        println!(
            "⚠ path A diverges from path B by >10× the pure-quant noise — qmatmul kernel BUG suspected"
        );
    } else if a_vs_b_rms < 1e-4 {
        println!("✓ path A and path B agree within rounding — qmatmul kernel CORRECT");
    }
    println!("first 8 values:");
    println!("  A: {:?}", &a_vals[..8.min(a_vals.len())]);
    println!("  B: {:?}", &b_vals[..8.min(b_vals.len())]);
    println!("  C: {:?}", &c_vals[..8.min(c_vals.len())]);
    Ok(())
}

fn main() -> Result<()> {
    let backend = std::env::args().nth(1).unwrap_or_else(|| "cpu".into());
    let device = match backend.as_str() {
        "metal" => Device::new_metal(0)?,
        _ => Device::Cpu,
    };
    println!("device: {device:?}");

    // 2D inputs (M, K) so both matmul ranks line up; QMatMul handles batched
    // inputs internally.
    let weight_small = Tensor::randn(0.0f32, 0.02, (32, 32), &device)?;
    let xs_small = Tensor::randn(0.0f32, 1.0, (4, 32), &device)?;
    run_case(
        "small 32x32",
        &weight_small,
        &xs_small,
        GgmlDType::Q8_0,
        &device,
    )?;

    let weight_attn = Tensor::randn(0.0f32, 0.02, (2048, 2048), &device)?;
    let xs_attn = Tensor::randn(0.0f32, 1.0, (8, 2048), &device)?;
    run_case(
        "attn 2048x2048",
        &weight_attn,
        &xs_attn,
        GgmlDType::Q8_0,
        &device,
    )?;

    let weight_kv = Tensor::randn(0.0f32, 0.02, (512, 2048), &device)?;
    let xs_kv = Tensor::randn(0.0f32, 1.0, (8, 2048), &device)?;
    run_case("kv 512x2048", &weight_kv, &xs_kv, GgmlDType::Q8_0, &device)?;

    run_case(
        "attn 2048x2048 (Q4_K)",
        &weight_attn,
        &xs_attn,
        GgmlDType::Q4K,
        &device,
    )?;

    // GGUF round-trip integrity test: write a real CSM weight to GGUF, read
    // it back, see if the values match.
    println!("\n=== GGUF round-trip integrity ===");
    use candle_core::quantized::gguf_file;
    let tmp = std::env::temp_dir().join("csm_qrt_test.gguf");
    let qt_orig = QTensor::quantize(&weight_attn, GgmlDType::Q8_0)?;
    let deq_orig = qt_orig.dequantize(&device)?.to_dtype(DType::F32)?;
    {
        let f = std::fs::File::create(&tmp)?;
        let mut w = std::io::BufWriter::new(f);
        let metadata: Vec<(&str, &gguf_file::Value)> = vec![];
        let tensors: Vec<(&str, &QTensor)> = vec![("test.weight", &qt_orig)];
        gguf_file::write(&mut w, &metadata, &tensors)?;
    }
    println!("wrote {} bytes", std::fs::metadata(&tmp)?.len());

    let mut f = std::fs::File::open(&tmp)?;
    let ct = gguf_file::Content::read(&mut f)?;
    let qt_round = ct.tensor(&mut f, "test.weight", &device)?;
    println!(
        "round-trip qt shape={:?} dtype={:?}",
        qt_round.shape(),
        qt_round.dtype()
    );
    let deq_round = qt_round.dequantize(&device)?.to_dtype(DType::F32)?;

    let v_orig = deq_orig.flatten_all()?.to_vec1::<f32>()?;
    let v_round = deq_round.flatten_all()?.to_vec1::<f32>()?;
    let r = rms(&v_orig, &v_round);
    let m = max_abs_diff(&v_orig, &v_round);
    println!("dequant(orig) vs dequant(round-trip): rms={r:.8}  max={m:.8}");
    if r > 1e-6 {
        println!("⚠ GGUF round-trip CORRUPTS data");
    } else {
        println!("✓ GGUF round-trip preserves data");
    }
    let _ = std::fs::remove_file(&tmp);

    // Now test qmatmul on the round-tripped tensor.
    let qmm_round = QMatMul::from_qtensor(qt_round)?;
    let out_round = qmm_round.forward(&xs_attn)?;
    let qt_orig_b = QTensor::quantize(&weight_attn, GgmlDType::Q8_0)?;
    let qmm_orig = QMatMul::from_qtensor(qt_orig_b)?;
    let out_orig = qmm_orig.forward(&xs_attn)?;
    let v_o = out_orig.flatten_all()?.to_vec1::<f32>()?;
    let v_r = out_round.flatten_all()?.to_vec1::<f32>()?;
    let r2 = rms(&v_o, &v_r);
    println!("qmatmul(orig) vs qmatmul(round-trip): rms={r2:.8}");

    Ok(())
}