//! Generator of `src/cohere2_moe_reference.txt`, the stored reference
//! the test `cohere2_moe::tests::stored_reference` checks.
//!
//! It builds the test's three synthetic Cohere2-MoE checkpoints (the
//! same weights, byte for byte), runs the test's scenario through
//! `forward_prompt` and `forward_token`, and prints one line per output
//! in the file's format, plus a `math-library` line naming this
//! system's math library, on which the bits depend as well as on the
//! code.
//!
//! The stored lines are this program's output at commit 30c4cf1ee (main
//! before the CPU passes), whose `forward_prompt` and `forward_token`
//! are the code the CPU passes must match. It uses only API that commit
//! has, so to add a platform:
//!
//! ```bash
//! git worktree add --detach /tmp/base 30c4cf1ee
//! cp examples/cohere2_reference.rs /tmp/base/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/
//! cd /tmp/base/jouleclaw/jouleclaw-rs
//! cargo run -q -p jouleclaw-loader-gguf --example cohere2_reference
//! cargo run -q -p jouleclaw-loader-gguf --features rayon --example cohere2_reference   # must print the same
//! ```
//!
//! and append its lines to the file. Run at a later commit, it prints
//! that commit's lines, which the test requires to be the same.

use jouleclaw_loader_gguf::cohere2_moe::{Cohere2MoeConfig, Cohere2MoeModel};
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;

fn tiny_config() -> Cohere2MoeConfig {
    Cohere2MoeConfig {
        hidden_size: 8,
        num_hidden_layers: 4,
        vocab_size: 13,
        layer_types: vec![
            "full_attention".into(),
            "sliding_attention".into(),
            "sliding_attention".into(),
            "full_attention".into(),
        ],
        num_attention_heads: 2,
        num_key_value_heads: 1,
        head_dim: 4,
        rope_theta: 10_000.0,
        sliding_window: 3,
        num_experts: 4,
        num_experts_per_tok: 2,
        intermediate_size: 6,
        expert_selection_fn: "sigmoid".into(),
        norm_topk_prob: false,
        first_k_dense_replace: 1,
        prefix_dense_intermediate_size: Some(10),
        prefix_dense_sliding_window_pattern: 1,
        rms_norm_eps: 1e-6,
        logit_scale: 0.5,
    }
}

fn medium_config() -> Cohere2MoeConfig {
    Cohere2MoeConfig {
        hidden_size: 64,
        num_hidden_layers: 4,
        vocab_size: 37,
        layer_types: vec![
            "full_attention".into(),
            "sliding_attention".into(),
            "sliding_attention".into(),
            "full_attention".into(),
        ],
        num_attention_heads: 4,
        num_key_value_heads: 2,
        head_dim: 32,
        rope_theta: 10_000.0,
        sliding_window: 20,
        num_experts: 8,
        num_experts_per_tok: 3,
        intermediate_size: 24,
        expert_selection_fn: "sigmoid".into(),
        norm_topk_prob: false,
        first_k_dense_replace: 1,
        prefix_dense_intermediate_size: Some(40),
        prefix_dense_sliding_window_pattern: 1,
        rms_norm_eps: 1e-6,
        logit_scale: 0.5,
    }
}

fn wide_config() -> Cohere2MoeConfig {
    Cohere2MoeConfig {
        hidden_size: 256,
        num_hidden_layers: 3,
        vocab_size: 1031,
        layer_types: vec![
            "full_attention".into(),
            "sliding_attention".into(),
            "full_attention".into(),
        ],
        num_attention_heads: 8,
        num_key_value_heads: 2,
        head_dim: 32,
        rope_theta: 10_000.0,
        sliding_window: 16,
        num_experts: 8,
        num_experts_per_tok: 2,
        intermediate_size: 96,
        expert_selection_fn: "sigmoid".into(),
        norm_topk_prob: false,
        first_k_dense_replace: 1,
        prefix_dense_intermediate_size: Some(2048),
        prefix_dense_sliding_window_pattern: 1,
        rms_norm_eps: 1e-6,
        logit_scale: 0.5,
    }
}

/// A scenario: `prompt` tokens prefilled, `decode` steps, `extend`
/// tokens fed into the live state, two more steps.
struct Shape {
    prompt: u32,
    decode: u32,
    extend: u32,
}

/// Deterministic weights for a tensor, from its name.
fn fill(name: &str, n: usize) -> Vec<f32> {
    let mut h = 1469598103934665603u64;
    for b in name.bytes() {
        h = (h ^ b as u64).wrapping_mul(1099511628211);
    }
    (0..n)
        .map(|i| {
            let mut x = h.wrapping_add((i as u64).wrapping_mul(0x9E3779B97F4A7C15));
            x ^= x >> 33;
            x = x.wrapping_mul(0xFF51AFD7ED558CCD);
            x ^= x >> 33;
            ((x % 2000) as f32 / 1000.0 - 1.0) * 0.4
        })
        .collect()
}

/// The checkpoint for `cfg` as a BF16 safetensors file at `path`.
fn build_checkpoint(path: &std::path::Path, cfg: &Cohere2MoeConfig) {
    let h = cfg.hidden_size as u64;
    let mut tensors: Vec<(String, Vec<u64>, Vec<f32>)> = Vec::new();
    let mut push = |name: String, shape: Vec<u64>| {
        let n: u64 = shape.iter().product();
        let vals = fill(&name, n as usize);
        tensors.push((name, shape, vals));
    };
    push("model.embed_tokens.weight".into(), vec![cfg.vocab_size as u64, h]);
    push("model.norm.weight".into(), vec![h]);
    let q_dim = (cfg.num_attention_heads * cfg.head_dim) as u64;
    let kv_dim = (cfg.num_key_value_heads * cfg.head_dim) as u64;
    for i in 0..cfg.num_hidden_layers {
        let p = format!("model.layers.{i}");
        push(format!("{p}.input_layernorm.weight"), vec![h]);
        push(format!("{p}.self_attn.q_proj.weight"), vec![q_dim, h]);
        push(format!("{p}.self_attn.k_proj.weight"), vec![kv_dim, h]);
        push(format!("{p}.self_attn.v_proj.weight"), vec![kv_dim, h]);
        push(format!("{p}.self_attn.o_proj.weight"), vec![h, q_dim]);
        if i < cfg.first_k_dense_replace {
            let pf = cfg.prefix_dense_intermediate_size.unwrap() as u64;
            push(format!("{p}.mlp.gate_proj.weight"), vec![pf, h]);
            push(format!("{p}.mlp.up_proj.weight"), vec![pf, h]);
            push(format!("{p}.mlp.down_proj.weight"), vec![h, pf]);
        } else {
            push(format!("{p}.mlp.gate.weight"), vec![cfg.num_experts as u64, h]);
            for e in 0..cfg.num_experts {
                let f = cfg.intermediate_size as u64;
                push(format!("{p}.mlp.experts.{e}.gate_proj.weight"), vec![f, h]);
                push(format!("{p}.mlp.experts.{e}.up_proj.weight"), vec![f, h]);
                push(format!("{p}.mlp.experts.{e}.down_proj.weight"), vec![h, f]);
            }
        }
    }
    let mut header = String::from("{");
    let mut data: Vec<u8> = Vec::new();
    for (i, (name, shape, vals)) in tensors.iter().enumerate() {
        let start = data.len();
        for &v in vals {
            // BF16 by truncation, as the test's checkpoints are.
            data.extend_from_slice(&((v.to_bits() >> 16) as u16).to_le_bytes());
        }
        let end = data.len();
        if i > 0 {
            header.push(',');
        }
        header.push_str(&format!(
            "\"{name}\":{{\"dtype\":\"BF16\",\"shape\":{shape:?},\"data_offsets\":[{start},{end}]}}"
        ));
    }
    header.push('}');
    let mut out = Vec::new();
    out.extend_from_slice(&(header.len() as u64).to_le_bytes());
    out.extend_from_slice(header.as_bytes());
    out.extend_from_slice(&data);
    std::fs::write(path, out).expect("write the checkpoint");
}

/// FNV-1a over every logit's bits, little-endian (north_bench's hash).
fn logits_hash(logits: &[f32]) -> u64 {
    let mut h = 0xcbf2_9ce4_8422_2325u64;
    for v in logits {
        for b in v.to_bits().to_le_bytes() {
            h = (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3);
        }
    }
    h
}

/// The scenario's lines for one model: `<model> <output> <hash>`.
fn scenario(name: &str, model: &Cohere2MoeModel, shape: &Shape) -> Vec<String> {
    let v = model.config.vocab_size as u32;
    let mut out = Vec::new();
    let mut line = |what: String, l: &[f32]| out.push(format!("{name} {what} {:016x}", logits_hash(l)));
    let mut state = model.new_state();
    let prompt: Vec<u32> = (0..shape.prompt).map(|i| (i * 7 + 3) % v).collect();
    line("prefill".into(), &model.forward_prompt(&prompt, &mut state));
    for i in 0..shape.decode {
        line(format!("decode{i}"), &model.forward_token((i * 5 + 2) % v, &mut state));
    }
    // The chunk is fed one decode step at a time (`forward_extend`,
    // which batches it, gives the same logits and does not exist at
    // the commit the stored lines come from).
    let mut fed = Vec::new();
    for i in 0..shape.extend {
        fed = model.forward_token((i * 11 + 5) % v, &mut state);
    }
    line("extend".into(), &fed);
    for i in 0..2 {
        line(format!("after{i}"), &model.forward_token((i * 3 + 1) % v, &mut state));
    }
    out
}

/// This system's math library, as `math-library` names it.
fn math_library() -> String {
    let run = |cmd: &str, args: &[&str]| {
        std::process::Command::new(cmd)
            .args(args)
            .output()
            .ok()
            .map(|o| String::from_utf8_lossy(&o.stdout).lines().next().unwrap_or("").trim().to_string())
            .unwrap_or_default()
    };
    match std::env::consts::OS {
        "macos" => format!(
            "the system math library (libsystem_m) of macOS {} ({})",
            run("sw_vers", &["-productVersion"]),
            run("sw_vers", &["-buildVersion"])
        ),
        "linux" => {
            // `ldd (Ubuntu GLIBC 2.39-0ubuntu8.9) 2.39`: the C library's name and version.
            let l = run("ldd", &["--version"]);
            let libc = l.split_once('(').and_then(|(_, r)| r.rsplit_once(')')).map_or(l.as_str(), |(n, _)| n);
            format!("the libm of {libc}")
        }
        os => format!("the system math library of {os}"),
    }
}

fn main() {
    let platform = format!("{}-{}", std::env::consts::OS, std::env::consts::ARCH);
    println!("math-library {platform} {}", math_library());
    let models = [
        ("tiny", tiny_config(), Shape { prompt: 5, decode: 4, extend: 4 }),
        ("medium", medium_config(), Shape { prompt: 45, decode: 6, extend: 25 }),
        ("wide", wide_config(), Shape { prompt: 24, decode: 4, extend: 20 }),
    ];
    for (name, cfg, shape) in models {
        let path = std::env::temp_dir().join(format!("cohere2-reference-{name}-{}.safetensors", std::process::id()));
        build_checkpoint(&path, &cfg);
        let st = read_safetensors_file(&path).expect("read the checkpoint");
        let model = Cohere2MoeModel::bind(cfg, st).expect("bind the checkpoint");
        let _ = std::fs::remove_file(&path);
        for line in scenario(name, &model, &shape) {
            println!("{platform} {line}");
        }
    }
}