//! End-to-end training entry point (Phase T6): load the GPT-2 BPE + TinyStories
//! corpus, train the tiny transformer with hand-written AdamW for a BOUNDED
//! budget, checkpoint it, and print a few generated samples.
//!
//! Run on dash5 (needs a GPU + the corpus + tokenizer.json):
//!   export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
//!   cargo run -p xtrain-train --release --bin train -- \
//!       /opt/wjh/models/gpt2/tokenizer.json \
//!       data/tinystories-valid-3mb.txt
//!
//! Optional 3rd/4th args: number of steps, checkpoint path.

// On a GPU-less host (no_cuda) the whole training body is unavailable; keep a
// stub `main` so the crate still builds for `cargo check`.
#[cfg(no_cuda)]
fn main() {
    eprintln!("xtrain train: built without CUDA (no_cuda); run on a GPU host (dash5).");
}

#[cfg(not(no_cuda))]
use std::path::{Path, PathBuf};

#[cfg(not(no_cuda))]
use xtrain_cuda::device;
#[cfg(not(no_cuda))]
use xtrain_model::{Config, TinyTransformer};
#[cfg(not(no_cuda))]
use xtrain_tensor::Device;
#[cfg(not(no_cuda))]
use xtrain_train::data::Corpus;
#[cfg(not(no_cuda))]
use xtrain_train::sample::generate;
#[cfg(not(no_cuda))]
use xtrain_train::schedule::LrSchedule;
#[cfg(not(no_cuda))]
use xtrain_train::{TrainConfig, train};

// Deterministic LCG fill in [-scale, scale) — same init scheme as the T5 tests.
#[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 main() {
    let args: Vec<String> = std::env::args().collect();
    let tok_path = args
        .get(1)
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("/opt/wjh/models/gpt2/tokenizer.json"));
    let corpus_path = args
        .get(2)
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("data/tinystories-valid-3mb.txt"));
    let steps: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(2000);
    let ckpt: PathBuf = args
        .get(4)
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("/tmp/xtrain_tinystories.ckpt"));

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

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

    // Tiny model sized to the BPE vocab. A real (but small) config: wider than
    // the overfit test so it has capacity to learn English structure.
    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 mut seed = 1u64;
    let model = TinyTransformer::new(cfg, device, |shape| {
        seed = seed.wrapping_add(1);
        let n: usize = shape.iter().product();
        if shape.len() == 1 {
            // RMSNorm gammas → ~1.
            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
        } else {
            // Small fan-in-ish scale; keeps early logits tame.
            fill(n, seed, 0.04)
        }
    });

    let seq_len = 64;
    let tcfg = TrainConfig {
        seq_len,
        batch_size: 8,
        steps,
        schedule: LrSchedule {
            max_lr: 3e-3,
            min_lr: 3e-4,
            warmup: (steps / 20).max(20),
            total: steps,
        },
        weight_decay: 0.1,
        max_grad_norm: 1.0,
        log_every: 50,
        ckpt_path: Some(ckpt.clone()),
        ckpt_every: 500,
        seed: 42,
    };

    println!(
        "training: {} steps, seq {}, batch {}, lr {:.1e}→{:.1e}",
        tcfg.steps, tcfg.seq_len, tcfg.batch_size, tcfg.schedule.max_lr, tcfg.schedule.min_lr
    );
    let losses = train(&model, device, &corpus, &tcfg);
    let start = losses.first().copied().unwrap_or(0.0);
    let end = losses.last().copied().unwrap_or(0.0);
    println!("loss: start {start:.4} → end {end:.4}");

    sample_some(&model, device, &tok_path);
}

#[cfg(not(no_cuda))]
fn sample_some(model: &TinyTransformer, device: Device, tok_path: &Path) {
    use xserv_tokenizer::Tokenizer;
    let tok = Tokenizer::from_file(tok_path);
    let prompts = ["Once upon a time", "The little", "One day"];
    println!("\n--- samples (greedy) ---");
    for p in prompts {
        let ids: Vec<i32> = tok.encode(p).into_iter().map(|t| t as i32).collect();
        let mut rng = 7u64;
        let out = generate(model, device, &ids, 40, 0.0, &mut rng);
        let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
        println!("[{p}] → {text}");
    }
    println!("\n--- samples (temperature 0.8) ---");
    for p in prompts {
        let ids: Vec<i32> = tok.encode(p).into_iter().map(|t| t as i32).collect();
        let mut rng = 13u64;
        let out = generate(model, device, &ids, 40, 0.8, &mut rng);
        let text = tok.decode(&out.iter().map(|&t| t as u32).collect::<Vec<_>>());
        println!("[{p}] → {text}");
    }
}
