//! Multi-rank DDP training launcher (Phase T8): spawn one thread per GPU, NCCL
//! all-reduce the gradients each step, and train the tiny transformer on
//! TinyStories. Doubles as the throughput driver — run it with 1/2/4 GPUs and
//! read the global tok/s line.
//!
//! Run on dash5 (pick idle GPUs — dash5 is shared):
//!   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//!   CUDA_VISIBLE_DEVICES=0,1 cargo run -p xtrain-distributed --release \
//!       --bin train_ddp -- 100 64 16
//! Args: [steps] [seq_len] [global_batch] [tokenizer.json] [corpus.txt]
//! The launcher uses every GPU visible to it (CUDA_VISIBLE_DEVICES selects them),
//! so the rank devices are always 0..N within the visible set.

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

#[cfg(not(no_cuda))]
fn main() {
    use std::path::PathBuf;
    use xtrain_cuda::device;
    use xtrain_distributed::{DdpConfig, build_model, launch};
    use xtrain_model::Config;
    use xtrain_train::data::Corpus;
    use xtrain_train::schedule::LrSchedule;

    let args: Vec<String> = std::env::args().collect();
    let steps: usize = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(100);
    let seq_len: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(64);
    let batch: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(16);
    let tok_path = args
        .get(4)
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
    let corpus_path = args
        .get(5)
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));

    // Use every visible GPU as a rank (CUDA_VISIBLE_DEVICES selects the set;
    // device ordinals are 0..count within it).
    let count = device::device_count().expect("device_count") as u32;
    assert!(count > 0, "no CUDA device visible");
    let devices: Vec<u32> = (0..count).collect();
    assert_eq!(
        batch % devices.len(),
        0,
        "global batch {batch} not divisible by world {}",
        devices.len()
    );

    println!(
        "DDP: world={} devices={:?} | steps={steps} seq={seq_len} global_batch={batch}",
        devices.len(),
        devices
    );

    let corpus = Corpus::load(&tok_path, &corpus_path);
    println!(
        "corpus: {} tokens, vocab {}",
        corpus.len(),
        corpus.vocab_size
    );

    let mut cfg = Config::tiny();
    cfg.vocab = corpus.vocab_size;
    cfg.n_layers = 4;
    println!(
        "model: dim {} layers {} heads {} ffn {} → {} params",
        cfg.dim,
        cfg.n_layers,
        cfg.n_heads,
        cfg.ffn_hidden,
        cfg.num_params()
    );

    let dcfg = DdpConfig {
        seq_len,
        batch_size: batch,
        steps,
        schedule: LrSchedule {
            max_lr: 3e-3,
            min_lr: 3e-4,
            warmup: (steps / 20).max(5),
            total: steps,
        },
        weight_decay: 0.1,
        max_grad_norm: 1.0,
        log_every: 10,
        seed: 42,
    };

    let traces = launch(&devices, &corpus, &dcfg, move |device| {
        build_model(cfg, device)
    });
    let trace = &traces[0];
    let start = trace.first().copied().unwrap_or(0.0);
    let end = trace.last().copied().unwrap_or(0.0);
    println!("loss: start {start:.4} → end {end:.4}");
}
