//! Dump gpt-oss next-token logits for a fixed token-id sequence, to compare
//! against the llama.cpp oracle (isolates the model forward from tokenizer
//! differences). Usage: gptoss-logits <mxfp4-model-dir> <tok0> <tok1> ...
use std::path::PathBuf;
use half::bf16;
use xserv_model::loader;
use xserv_model::{GptOss, ModelConfig};
use xserv_tensor::Device;

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let model_dir = PathBuf::from(&args[1]);
    let tokens: Vec<u32> = args[2..].iter().map(|s| s.parse().expect("token id")).collect();
    assert!(!tokens.is_empty(), "need at least one token id");

    xserv_cuda::device::set_device(0).unwrap();
    let config = ModelConfig::from_file(&model_dir.join("config.json"));
    eprintln!("[gptoss-logits] loading {} (MXFP4) ...", model_dir.display());
    let (floats, u8s) = loader::load_model_dir_split(&model_dir, Device::Cpu);
    let model = GptOss::from_weights(config, floats, u8s);
    eprintln!("[gptoss-logits] forward over {} tokens", tokens.len());

    // (1) batched host-attention forward (reference path).
    let logits = model.forward(&tokens); // [T, vocab]
    let vocab = logits.shape()[1];
    let t = logits.shape()[0];
    let host = logits.to_device(Device::Cpu);
    let data = host.as_slice::<bf16>();
    let last = &data[(t - 1) * vocab..t * vocab];
    let mut idx: Vec<usize> = (0..vocab).collect();
    idx.sort_by(|&a, &b| last[b].to_f32().partial_cmp(&last[a].to_f32()).unwrap());
    println!("[forward] top5 next-token (id: logit):");
    for &i in &idx[..5] {
        println!("  {i}: {:.4}", last[i].to_f32());
    }

    // (2) KV-cache GPU decode path (token-by-token prefill) — must match top-1.
    let mut cache = xserv_model::GpuKVCache::new(&model.config, 512, xserv_tensor::DType::BF16, 0);
    let mut dlog = model.decode_step(tokens[0], &mut cache);
    for &tok in &tokens[1..] {
        dlog = model.decode_step(tok, &mut cache);
    }
    let dh = dlog.to_device(Device::Cpu);
    let dd = dh.as_slice::<bf16>();
    let mut didx: Vec<usize> = (0..vocab).collect();
    didx.sort_by(|&a, &b| dd[b].to_f32().partial_cmp(&dd[a].to_f32()).unwrap());
    println!("[decode ] top5 next-token (id: logit):");
    for &i in &didx[..5] {
        println!("  {i}: {:.4}", dd[i].to_f32());
    }
    println!("MATCH_TOP1: {}", if idx[0] == didx[0] { "YES" } else { "NO" });
}