//! The training loop: sample sequences → forward `loss` → backward → grad clip
//! (with batch averaging) → AdamW step → zero grads; with an LR schedule,
//! periodic loss logging, and periodic checkpointing.
//!
//! The T5 model is single-sequence, so a "batch" of `batch_size` sequences is
//! handled by running forward+backward on each and letting the tape SUM their
//! grads (its fan-out rule); the clip pass then multiplies by `1/batch_size` to
//! recover the batch-mean gradient before clipping + the optimizer step.

#![cfg(not(no_cuda))]

use std::path::PathBuf;
use std::time::Instant;

use xtrain_model::{TinyTransformer, ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;

use crate::checkpoint;
use crate::clip::clip_grad_norm_gpu;
use crate::data::Corpus;
use crate::schedule::LrSchedule;

/// Knobs for a training run.
pub struct TrainConfig {
    pub seq_len: usize,
    pub batch_size: usize,
    pub steps: usize,
    pub schedule: LrSchedule,
    pub weight_decay: f32,
    pub max_grad_norm: f32,
    pub log_every: usize,
    /// Optional checkpoint path written every `ckpt_every` steps (and at the end).
    pub ckpt_path: Option<PathBuf>,
    pub ckpt_every: usize,
    /// Seed for reproducible sequence sampling.
    pub seed: u64,
}

/// Train `model` on `corpus` for `cfg.steps` AdamW steps. Returns the per-step
/// loss trace (one mean loss per step, read from the first sequence of the
/// batch — cheap and representative). Logs progress and checkpoints as configured.
pub fn train(
    model: &TinyTransformer,
    device: Device,
    corpus: &Corpus,
    cfg: &TrainConfig,
) -> Vec<f32> {
    let params = model.params();
    let mut opt = GpuAdamW::new(cfg.weight_decay);
    let mut rng = cfg.seed;
    let mut losses = Vec::with_capacity(cfg.steps);
    let inv_batch = 1.0 / cfg.batch_size as f32;
    let start = Instant::now();
    let mut tokens_seen: u64 = 0;

    for step in 0..cfg.steps {
        let lr = cfg.schedule.lr(step);

        // Accumulate grads over `batch_size` sequences (tape SUMs them).
        let mut step_loss = 0.0f32;
        for _ in 0..cfg.batch_size {
            let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
            let ids = ids_tensor(&input, device);
            let targets = ids_tensor(&target, device);
            let loss = model.loss(&ids, &targets);
            step_loss += read_scalar(&loss);
            loss.backward();
            tokens_seen += cfg.seq_len as u64;
        }
        step_loss *= inv_batch;
        losses.push(step_loss);

        // Average the summed grads (×1/batch) and clip to the global norm.
        let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch);
        opt.step(lr, &params);
        for p in &params {
            p.zero_grad();
        }

        if step % cfg.log_every == 0 || step == cfg.steps - 1 {
            let elapsed = start.elapsed().as_secs_f32();
            let tps = tokens_seen as f32 / elapsed.max(1e-6);
            println!(
                "step {step:5}/{}: loss {step_loss:.4} lr {lr:.2e} gnorm {gnorm:.3} \
                 ({tps:.0} tok/s)",
                cfg.steps
            );
        }

        if let Some(path) = &cfg.ckpt_path {
            if cfg.ckpt_every > 0 && (step + 1) % cfg.ckpt_every == 0 {
                checkpoint::save(path, &params).expect("checkpoint save");
            }
        }
    }

    if let Some(path) = &cfg.ckpt_path {
        checkpoint::save(path, &params).expect("final checkpoint save");
        println!("saved checkpoint → {}", path.display());
    }
    losses
}

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