//! Autoregressive text sampling from the trained model. The model is
//! single-sequence with RoPE position = row index, so generation re-runs the
//! forward on the growing prefix each step and reads the last row's logits — the
//! simplest correct approach (no KV cache; that is an inference/perf concern).
//!
//! Greedy when `temperature == 0`, else temperature sampling over the softmax.

#![cfg(not(no_cuda))]

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

/// Generate `max_new` tokens continuing `prompt`. `temperature == 0` → greedy
/// argmax; otherwise sample from softmax(logits / temperature). `rng_state` is a
/// reproducible LCG seed (only used when temperature > 0).
pub fn generate(
    model: &TinyTransformer,
    device: Device,
    prompt: &[i32],
    max_new: usize,
    temperature: f32,
    rng_state: &mut u64,
) -> Vec<i32> {
    let vocab = model.config().vocab;
    let mut ids: Vec<i32> = prompt.to_vec();

    for _ in 0..max_new {
        let ids_t = ids_tensor(&ids, device);
        // In bf16 mode the logits are bf16; cast to f32 (on device) before reading.
        let logits = model.forward(&ids_t).value();
        let logits = logits
            .to_dtype(xtrain_tensor::DType::F32)
            .to_device(Device::Cpu);
        let lg = logits.as_slice::<f32>();
        // Last row = next-token distribution for the current prefix.
        let last = &lg[(ids.len() - 1) * vocab..ids.len() * vocab];

        let next = if temperature <= 0.0 {
            argmax(last)
        } else {
            sample_temperature(last, temperature, rng_state)
        };
        ids.push(next as i32);
    }
    ids
}

fn argmax(row: &[f32]) -> usize {
    row.iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
        .unwrap()
        .0
}

fn sample_temperature(row: &[f32], temperature: f32, rng_state: &mut u64) -> usize {
    // Softmax with temperature (numerically stable).
    let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = row
        .iter()
        .map(|&x| ((x - max) / temperature).exp())
        .collect();
    let sum: f32 = exps.iter().sum();
    let r = (next_rand(rng_state) as f32 / u32::MAX as f32) * sum;
    let mut acc = 0.0;
    for (i, &e) in exps.iter().enumerate() {
        acc += e;
        if acc >= r {
            return i;
        }
    }
    exps.len() - 1
}

fn next_rand(state: &mut u64) -> u32 {
    *state = state
        .wrapping_mul(6364136223846793005)
        .wrapping_add(1442695040888963407);
    (*state >> 32) as u32
}
