//! DPO training on the verifiable arithmetic task (M3 / Stage P1).
//!
//! Loads the SFT checkpoint as the policy AND uses it as the frozen reference:
//! reference logprobs `log πref(chosen)` / `log πref(rejected)` are **precomputed
//! once** before any optimizer step (when policy == reference), then cached as
//! constants — so only one model stays resident (the design's reference-logprob
//! caching). Each step forwards the policy on the chosen and rejected completions,
//! takes [`seq_logprob`] of each, and minimises [`dpo_loss`]; the two forwards
//! share the policy params, so backward accumulates both branches' grads.
//!
//! Health metrics (per docs/18, the doc-13 "don't trust loss alone" lesson): the
//! chosen−rejected **reward margin** and **preference accuracy** (margin > 0) — both
//! should rise. The arithmetic-correctness payoff is measured separately by running
//! `eval_arith` on the saved checkpoint.
//!
//!   train_dpo <tokenizer.json> <dpo.tsv> --init-ckpt <sft.ckpt> <arch flags> \
//!       --beta 0.1 --steps 1000 --lr 5e-7 --ckpt <out.ckpt>

#[cfg(no_cuda)]
fn main() {
    eprintln!("train_dpo: built without CUDA (no_cuda); run on a GPU host.");
}

#[cfg(not(no_cuda))]
use xtrain_autodiff::ops;
#[cfg(not(no_cuda))]
use xtrain_cuda::device;
#[cfg(not(no_cuda))]
use xtrain_model::{Config, TinyTransformer, ids_tensor};
#[cfg(not(no_cuda))]
use xtrain_tensor::Device;

#[cfg(not(no_cuda))]
fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
    let mut state = seed
        .wrapping_mul(2862933555777941757)
        .wrapping_add(3037000493);
    (0..n)
        .map(|_| {
            state = state
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            (((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
        })
        .collect()
}

#[cfg(not(no_cuda))]
fn flag<T: std::str::FromStr>(args: &[String], name: &str, default: T) -> T {
    args.iter()
        .position(|a| a == name)
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(default)
}

#[cfg(not(no_cuda))]
fn flag_value(args: &[String], name: &str) -> Option<String> {
    args.iter()
        .position(|a| a == name)
        .and_then(|i| args.get(i + 1))
        .cloned()
}

/// Frame a (question, completion) the same way the SFT loader does
/// (`User: …\nAssistant:` prompt + ` {completion}\n<|endoftext|>`), then return the
/// next-token (input, target) pair: input = tokens[..L-1], target = labels[1..L]
/// with the prompt positions masked to -100 (only completion tokens supervised).
#[cfg(not(no_cuda))]
fn frame(
    tok: &xserv_tokenizer::Tokenizer,
    question: &str,
    completion: &str,
) -> (Vec<i32>, Vec<i32>) {
    let prompt = format!("User: {question}\nAssistant:");
    let answer = format!(" {completion}\n<|endoftext|>");
    let p_ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
    let a_ids: Vec<i32> = tok.encode(&answer).into_iter().map(|t| t as i32).collect();
    let mut tokens = p_ids.clone();
    tokens.extend_from_slice(&a_ids);
    let mut labels = vec![-100i32; p_ids.len()];
    labels.extend_from_slice(&a_ids);
    let l = tokens.len();
    (tokens[..l - 1].to_vec(), labels[1..l].to_vec())
}

/// Sequence logprob `Σ log πθ(completion)` of a framed (input, target) pair.
#[cfg(not(no_cuda))]
fn seq_lp(
    model: &TinyTransformer,
    device: Device,
    input: &[i32],
    target: &[i32],
) -> xtrain_autodiff::tape::Var {
    let logits = model.forward(&ids_tensor(input, device));
    ops::seq_logprob(&logits, &ids_tensor(target, device))
}

#[cfg(not(no_cuda))]
fn scalar(v: &xtrain_autodiff::tape::Var) -> f32 {
    v.value().to_device(Device::Cpu).as_slice::<f32>()[0]
}

#[cfg(not(no_cuda))]
fn main() {
    use xserv_tokenizer::Tokenizer;
    use xtrain_optim::GpuAdamW;

    let args: Vec<String> = std::env::args().collect();
    let positionals: Vec<&String> = args[1..].iter().filter(|a| !a.starts_with("--")).collect();
    let tok_path = positionals.first().expect("usage: train_dpo <tokenizer.json> <dpo.tsv> [flags]");
    let tsv_path = positionals.get(1).expect("usage: train_dpo <tokenizer.json> <dpo.tsv> [flags]");

    let n_heads = flag(&args, "--heads", 52usize);
    let head_dim = flag(&args, "--head-dim", 32usize);
    let n_layers = flag(&args, "--layers", 22usize);
    let ffn = flag(&args, "--ffn", 6656usize);
    let kv_heads = flag(&args, "--kv-heads", n_heads);
    let beta: f32 = flag(&args, "--beta", 0.1);
    let steps: usize = flag(&args, "--steps", 1000);
    let lr: f32 = flag(&args, "--lr", 5e-7);
    let wd: f32 = flag(&args, "--wd", 0.0);
    let clip: f32 = flag(&args, "--clip", 1.0);
    let log_every: usize = flag(&args, "--log-every", 50);
    let init_ckpt = flag_value(&args, "--init-ckpt").expect("--init-ckpt <sft.ckpt> is required");
    let out_ckpt = flag_value(&args, "--ckpt").expect("--ckpt <out> is required");

    // Load preference pairs: question<TAB>chosen<TAB>rejected.
    let raw = std::fs::read_to_string(tsv_path).expect("read dpo tsv");
    let pairs: Vec<(String, String, String)> = raw
        .lines()
        .filter(|l| !l.trim().is_empty())
        .map(|l| {
            let mut it = l.splitn(3, '\t');
            let q = it.next().expect("question").to_string();
            let c = it.next().expect("chosen").to_string();
            let r = it.next().expect("rejected").to_string();
            (q, c, r)
        })
        .collect();
    assert!(!pairs.is_empty(), "no DPO pairs in {tsv_path}");

    assert!(device::device_count().unwrap() > 0, "no CUDA device");
    device::set_device(0).unwrap();
    let device = Device::Cuda(0);

    let tok = Tokenizer::from_file(std::path::Path::new(tok_path.as_str()));
    let cfg = Config::from_arch(tok.vocab_size(), n_heads, head_dim, n_layers, ffn)
        .with_kv_heads(kv_heads);
    let mut seed_init = 1u64;
    let model = TinyTransformer::new(cfg, device, |shape| {
        seed_init = seed_init.wrapping_add(1);
        let n: usize = shape.iter().product();
        if shape.len() == 1 {
            fill(n, seed_init, 0.02).iter().map(|v| v + 1.0).collect()
        } else {
            fill(n, seed_init, 0.04)
        }
    });
    xtrain_train::checkpoint::load_into(std::path::Path::new(&init_ckpt), &model.params())
        .expect("load SFT checkpoint");
    model.eval(); // DPO runs without dropout (deterministic logprobs)

    // Pre-tokenize every pair once.
    let framed: Vec<((Vec<i32>, Vec<i32>), (Vec<i32>, Vec<i32>))> = pairs
        .iter()
        .map(|(q, c, r)| (frame(&tok, q, c), frame(&tok, q, r)))
        .collect();

    // Reference logprobs: computed ONCE while policy == reference (SFT init), cached.
    println!("precomputing reference logprobs for {} pairs…", framed.len());
    let mut ref_c = Vec::with_capacity(framed.len());
    let mut ref_r = Vec::with_capacity(framed.len());
    for ((ci, ct), (ri, rt)) in &framed {
        ref_c.push(scalar(&seq_lp(&model, device, ci, ct)));
        ref_r.push(scalar(&seq_lp(&model, device, ri, rt)));
    }

    let params = model.params();
    let mut opt = GpuAdamW::new(wd);
    let n = framed.len();
    // A fixed shuffle (LCG-strided) so steps sweep the dataset without bias.
    let mut order: Vec<usize> = (0..n).collect();
    let mut s = 0x9E3779B97F4A7C15u64;
    for i in (1..n).rev() {
        s = s.wrapping_mul(6364136223846793005).wrapping_add(1);
        let j = (s >> 33) as usize % (i + 1);
        order.swap(i, j);
    }

    let start = std::time::Instant::now();
    let (mut win_loss, mut win_margin, mut win_acc) = (0f32, 0f32, 0usize);
    for step in 0..steps {
        let i = order[step % n];
        let ((ci, ct), (ri, rt)) = &framed[i];
        let lpc = seq_lp(&model, device, ci, ct);
        let lpr = seq_lp(&model, device, ri, rt);
        let (lpc_v, lpr_v) = (scalar(&lpc), scalar(&lpr));
        let margin = (lpc_v - ref_c[i]) - (lpr_v - ref_r[i]); // implicit reward margin
        let loss = ops::dpo_loss(&lpc, &lpr, ref_c[i], ref_r[i], beta);
        win_loss += scalar(&loss);
        win_margin += margin;
        win_acc += (margin > 0.0) as usize;

        loss.backward();
        let _ = xtrain_train::clip::clip_grad_norm_gpu(&params, clip, 1.0);
        opt.step(lr, &params);
        for p in &params {
            p.zero_grad();
        }

        if (step + 1) % log_every == 0 || step == steps - 1 {
            let w = log_every.min(step + 1) as f32;
            println!(
                "step {:5}/{steps}: loss {:.4} | reward-margin {:+.4} | pref-acc {:.1}% | {:.1}s",
                step + 1,
                win_loss / w,
                win_margin / w,
                100.0 * win_acc as f32 / w,
                start.elapsed().as_secs_f32(),
            );
            win_loss = 0.0;
            win_margin = 0.0;
            win_acc = 0;
        }
    }

    xtrain_train::checkpoint::save(std::path::Path::new(&out_ckpt), &params).expect("save ckpt");
    println!(
        "DPO done: {} pairs, {steps} steps, beta {beta}, lr {lr:.1e} → {out_ckpt}",
        framed.len()
    );
}