//! Trace layer-by-layer divergence between CPU and GPU
use realizar::gguf::{MappedGGUFModel, OwnedQKVWeights, OwnedQuantizedModel};
use realizar::quantize::{fused_q4k_parallel_matvec, fused_q6k_parallel_matvec};

const GGUF_TYPE_Q4_K: u32 = 12;
const GGUF_TYPE_Q6_K: u32 = 14;

fn rms_norm(input: &[f32], weight: &[f32], eps: f32) -> Vec<f32> {
    let n = input.len();
    let rms = (input.iter().map(|x| x * x).sum::<f32>() / n as f32 + eps).sqrt();
    input
        .iter()
        .zip(weight)
        .map(|(x, w)| (x / rms) * w)
        .collect()
}

fn fused_matmul(input: &[f32], data: &[u8], qtype: u32, in_dim: usize, out_dim: usize) -> Vec<f32> {
    match qtype {
        GGUF_TYPE_Q4_K => fused_q4k_parallel_matvec(data, input, in_dim, out_dim).unwrap(),
        GGUF_TYPE_Q6_K => fused_q6k_parallel_matvec(data, input, in_dim, out_dim).unwrap(),
        _ => panic!("Unsupported qtype: {}", qtype),
    }
}

fn cpu_layer_forward(
    hidden: &mut [f32],
    layer: &realizar::gguf::OwnedQuantizedLayer,
    num_heads: usize,
    num_kv_heads: usize,
    eps: f32,
) {
    let hidden_dim = hidden.len();
    let head_dim = hidden_dim / num_heads;
    let group_size = num_heads / num_kv_heads;

    // RMSNorm
    let normed = rms_norm(hidden, &layer.attn_norm_weight, eps);

    // Q, K, V projections
    let (q_weight, k_weight, v_weight) = match &layer.qkv_weight {
        OwnedQKVWeights::Separate { q, k, v } => (q, k, v),
        _ => panic!("Expected separate QKV"),
    };

    let _q = fused_matmul(
        &normed,
        &q_weight.data,
        q_weight.qtype,
        q_weight.in_dim,
        q_weight.out_dim,
    );
    let _k = fused_matmul(
        &normed,
        &k_weight.data,
        k_weight.qtype,
        k_weight.in_dim,
        k_weight.out_dim,
    );
    let v = fused_matmul(
        &normed,
        &v_weight.data,
        v_weight.qtype,
        v_weight.in_dim,
        v_weight.out_dim,
    );

    // Single-token attention: output = V (softmax(single element) = 1.0)
    let mut attn_out = Vec::with_capacity(hidden_dim);
    for h in 0..num_heads {
        let kv_head = h / group_size;
        let kv_start = kv_head * head_dim;
        attn_out.extend_from_slice(&v[kv_start..kv_start + head_dim]);
    }

    // Output projection
    let out_proj = fused_matmul(
        &attn_out,
        &layer.attn_output_weight.data,
        layer.attn_output_weight.qtype,
        layer.attn_output_weight.in_dim,
        layer.attn_output_weight.out_dim,
    );

    // Residual 1
    for (h, o) in hidden.iter_mut().zip(out_proj.iter()) {
        *h += o;
    }

    // FFN norm
    let ffn_normed = rms_norm(hidden, layer.ffn_norm_weight.as_ref().unwrap(), eps);

    // FFN gate + up
    let gate_weight = layer.ffn_gate_weight.as_ref().unwrap();
    let gate = fused_matmul(
        &ffn_normed,
        &gate_weight.data,
        gate_weight.qtype,
        gate_weight.in_dim,
        gate_weight.out_dim,
    );
    let up = fused_matmul(
        &ffn_normed,
        &layer.ffn_up_weight.data,
        layer.ffn_up_weight.qtype,
        layer.ffn_up_weight.in_dim,
        layer.ffn_up_weight.out_dim,
    );

    // SwiGLU
    let swiglu: Vec<f32> = gate
        .iter()
        .zip(up.iter())
        .map(|(g, u)| (g / (1.0 + (-*g).exp())) * u)
        .collect();

    // FFN down
    let ffn_down = fused_matmul(
        &swiglu,
        &layer.ffn_down_weight.data,
        layer.ffn_down_weight.qtype,
        layer.ffn_down_weight.in_dim,
        layer.ffn_down_weight.out_dim,
    );

    // Residual 2
    for (h, f) in hidden.iter_mut().zip(ffn_down.iter()) {
        *h += f;
    }
}

fn main() {
    let path = "/home/noah/models/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf";
    let mapped = MappedGGUFModel::from_path(path).unwrap();
    let model = OwnedQuantizedModel::from_mapped(&mapped).unwrap();

    let hidden_dim = model.config.hidden_dim;
    let num_heads = model.config.num_heads;
    let num_kv_heads = model.config.num_kv_heads;
    let num_layers = model.config.num_layers;
    let eps = model.config.eps;

    // Token 791
    let token_id = 791u32;
    let start = token_id as usize * hidden_dim;
    let mut hidden: Vec<f32> = model.token_embedding[start..start + hidden_dim].to_vec();

    println!("=== CPU Layer-by-Layer Forward ===");
    println!(
        "Token: {}, hidden_dim: {}, num_layers: {}",
        token_id, hidden_dim, num_layers
    );

    for layer_idx in 0..num_layers {
        cpu_layer_forward(
            &mut hidden,
            &model.layers[layer_idx],
            num_heads,
            num_kv_heads,
            eps,
        );
        println!(
            "Layer {:2}: first 3 = [{:.4}, {:.4}, {:.4}], sum={:.4}",
            layer_idx,
            hidden[0],
            hidden[1],
            hidden[2],
            hidden.iter().sum::<f32>()
        );
    }

    // Final RMSNorm + LM head
    let normed = rms_norm(&hidden, &model.output_norm_weight, eps);
    println!(
        "\nNormed hidden: first 3 = [{:.4}, {:.4}, {:.4}]",
        normed[0], normed[1], normed[2]
    );

    let logits = fused_matmul(
        &normed,
        &model.lm_head_weight.data,
        model.lm_head_weight.qtype,
        model.lm_head_weight.in_dim,
        model.lm_head_weight.out_dim,
    );

    let argmax = logits
        .iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
        .unwrap();
    println!("CPU argmax: {} (logit: {:.4})", argmax.0, argmax.1);
    println!("logit[16]: {:.4}", logits[16]);
    println!("logit[74403]: {:.4}", logits[74403]);
}
