//! MicroGPT Engine — 完整 Decoder-only GPT 实现 (Phase-68)
//!
//! 所有组件在单个文件中实现：
//! - T68.1: GptConfig + Token/Position Embedding
//! - T68.2: RmsNorm + softmax
//! - T68.3: CausalSelfAttention
//! - T68.4: Mlp + TransformerBlock
//! - T68.5: MicroGpt::forward() + DeviceSelector
//! - T68.6: train_step() + AdamW
//! - T68.7: generate() — 自回归温度采样
//! - T68.8: TokenizerManager — 字符级 + BPE

use candle_core::{DType, Device, IndexOp, Module, Result as CandleResult, Tensor, D};
use candle_nn::{
    embedding, linear, linear_no_bias, ops::softmax, Embedding, Linear, Optimizer, VarBuilder,
    VarMap,
};
use rand::Rng;
use serde::{Deserialize, Serialize};

use crate::device::DeviceSelector;

// ============================================================================
// T68.1 — GptConfig + Token/Position Embedding
// ============================================================================

/// GPT 模型配置
///
/// 定义了 MicroGPT 的所有超参数。默认值对应一个极小模型（~100K 参数），
/// 适合在 CPU 上快速训练验证。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GptConfig {
    /// 词表大小（字符级=256, BPE 可为 32K+）
    pub vocab_size: usize,
    /// 嵌入维度（隐藏层大小）
    pub n_embd: usize,
    /// 注意力头数
    pub n_head: usize,
    /// Transformer 层数
    pub n_layer: usize,
    /// 最大序列长度（上下文窗口）
    pub block_size: usize,
    /// 前馈网络中间层维度倍数（intermediate = n_embd * mlp_ratio）
    #[serde(default = "default_mlp_ratio")]
    pub mlp_ratio: usize,
    /// RMSNorm epsilon
    #[serde(default = "default_rms_eps")]
    pub rms_norm_eps: f64,
}

fn default_mlp_ratio() -> usize {
    4
}
fn default_rms_eps() -> f64 {
    1e-5
}

impl Default for GptConfig {
    fn default() -> Self {
        Self {
            vocab_size: 256, // 字符级
            n_embd: 64,
            n_head: 4,
            n_layer: 2,
            block_size: 128,
            mlp_ratio: 4,
            rms_norm_eps: 1e-5,
        }
    }
}

impl GptConfig {
    /// 创建一个微型配置（用于测试，~25K 参数）
    pub fn tiny() -> Self {
        Self {
            vocab_size: 256,
            n_embd: 32,
            n_head: 2,
            n_layer: 1,
            block_size: 64,
            mlp_ratio: 4,
            rms_norm_eps: 1e-5,
        }
    }

    /// 创建一个小型配置（~2M 参数）
    pub fn small() -> Self {
        Self {
            vocab_size: 256,
            n_embd: 128,
            n_head: 4,
            n_layer: 4,
            block_size: 256,
            mlp_ratio: 4,
            rms_norm_eps: 1e-5,
        }
    }

    /// 每个注意力头的维度
    pub fn head_dim(&self) -> usize {
        self.n_embd / self.n_head
    }

    /// 前馈网络中间层维度
    pub fn intermediate_size(&self) -> usize {
        self.n_embd * self.mlp_ratio
    }

    /// 估算模型参数量
    pub fn estimate_params(&self) -> usize {
        let embd = self.vocab_size * self.n_embd + self.block_size * self.n_embd;
        let attn_per_layer = 4 * self.n_embd * self.n_embd; // Q, K, V, O
        let mlp_per_layer = 2 * self.n_embd * self.intermediate_size();
        let norm_per_layer = 2 * self.n_embd; // 2x RMSNorm
        let lm_head = self.n_embd * self.vocab_size;
        let final_norm = self.n_embd;
        embd + self.n_layer * (attn_per_layer + mlp_per_layer + norm_per_layer)
            + final_norm
            + lm_head
    }
}

// ============================================================================
// T68.2 — RmsNorm + Softmax
// ============================================================================

/// Root Mean Square Layer Normalization
///
/// 比 LayerNorm 更高效，不需要计算均值。
/// 公式: `output = x / RMS(x) * weight`
/// 其中 `RMS(x) = sqrt(mean(x^2) + eps)`
struct RmsNorm {
    weight: Tensor,
    eps: f64,
}

impl RmsNorm {
    /// 从 VarBuilder 加载/初始化 RmsNorm
    fn new(dim: usize, eps: f64, vb: VarBuilder) -> CandleResult<Self> {
        let weight = vb.get_with_hints(dim, "weight", candle_nn::Init::Const(1.0))?;
        Ok(Self { weight, eps })
    }
}

impl Module for RmsNorm {
    fn forward(&self, x: &Tensor) -> CandleResult<Tensor> {
        let dtype = x.dtype();
        let x = x.to_dtype(DType::F32)?;
        let variance = x.sqr()?.mean_keepdim(D::Minus1)?;
        let rms = (variance + self.eps)?.sqrt()?;
        let normalized = x.broadcast_div(&rms)?;
        let weight = self.weight.to_dtype(DType::F32)?;
        normalized.broadcast_mul(&weight)?.to_dtype(dtype)
    }
}

/// 因果注意力掩码：下三角矩阵
///
/// 返回形状为 (1, 1, T, T) 的掩码，上三角部分为 -inf。
/// 使得 softmax 后上三角区域的注意力权重为 0。
pub(crate) fn causal_mask(seq_len: usize, device: &Device) -> CandleResult<Tensor> {
    let mask: Vec<f32> = (0..seq_len)
        .flat_map(|i| {
            (0..seq_len).map(move |j| {
                if j <= i {
                    0.0f32
                } else {
                    f32::NEG_INFINITY
                }
            })
        })
        .collect();
    Tensor::from_vec(mask, (1, 1, seq_len, seq_len), device)
}

// ============================================================================
// T68.3 — Multi-Head Causal Self-Attention
// ============================================================================

/// 多头因果自注意力
///
/// GPT 的核心组件。实现 Q·K^T/√d → causal_mask → softmax → V 的计算。
/// 每个 token 只能看到自己和之前的 token（因果掩码）。
struct CausalSelfAttention {
    /// QKV 合并投影 (n_embd → 3 * n_embd)，减少 kernel launch
    qkv_proj: Linear,
    /// 输出投影 (n_embd → n_embd)
    out_proj: Linear,
    /// 注意力头数
    n_head: usize,
    /// 每个头的维度
    head_dim: usize,
    /// 缩放因子 1/√d
    scale: f64,
}

impl CausalSelfAttention {
    fn new(config: &GptConfig, vb: VarBuilder) -> CandleResult<Self> {
        let n_embd = config.n_embd;
        let head_dim = config.head_dim();
        // QKV 合并投影：一次 matmul 计算 Q, K, V
        let qkv_proj = linear(n_embd, 3 * n_embd, vb.pp("qkv_proj"))?;
        let out_proj = linear(n_embd, n_embd, vb.pp("out_proj"))?;
        let scale = (head_dim as f64).powf(-0.5);

        Ok(Self {
            qkv_proj,
            out_proj,
            n_head: config.n_head,
            head_dim,
            scale,
        })
    }

    /// 前向传播
    ///
    /// x: (B, T, C) → output: (B, T, C)
    fn forward(&self, x: &Tensor, mask: &Tensor) -> CandleResult<Tensor> {
        let (b, t, c) = x.dims3()?;
        let n_embd = self.n_head * self.head_dim;

        // QKV 合并投影: (B, T, C) → (B, T, 3C)
        let qkv = self.qkv_proj.forward(x)?;

        // 拆分 Q, K, V: 各 (B, T, C)
        let q = qkv.narrow(D::Minus1, 0, n_embd)?;
        let k = qkv.narrow(D::Minus1, n_embd, n_embd)?;
        let v = qkv.narrow(D::Minus1, 2 * n_embd, n_embd)?;

        // 重塑为多头: (B, T, n_head, head_dim) → (B, n_head, T, head_dim)
        let q = q
            .reshape((b, t, self.n_head, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let k = k
            .reshape((b, t, self.n_head, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let v = v
            .reshape((b, t, self.n_head, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;

        // 注意力分数: Q @ K^T * scale → (B, n_head, T, T)
        let k_t = k.transpose(D::Minus2, D::Minus1)?;
        let attn = q.matmul(&k_t)?;
        let attn = (attn * self.scale)?;

        // 应用因果掩码: 上三角 → -inf → softmax 后为 0
        let attn = attn.broadcast_add(mask)?;
        let attn = softmax(&attn, D::Minus1)?;

        // 加权求和: (B, n_head, T, T) @ (B, n_head, T, head_dim) → (B, n_head, T, head_dim)
        let out = attn.matmul(&v.contiguous()?)?;

        // 合并多头: (B, n_head, T, head_dim) → (B, T, C)
        let out = out.transpose(1, 2)?.contiguous()?;
        let out = out.reshape((b, t, c))?;

        // 输出投影
        self.out_proj.forward(&out)
    }
}

// ============================================================================
// T68.4 — MLP Block + TransformerBlock
// ============================================================================

/// 前馈网络 (MLP)
///
/// 两层线性变换 + GELU 激活：
/// `output = down_proj(GELU(up_proj(x)))`
struct Mlp {
    up_proj: Linear,
    down_proj: Linear,
}

impl Mlp {
    fn new(config: &GptConfig, vb: VarBuilder) -> CandleResult<Self> {
        let n_embd = config.n_embd;
        let intermediate = config.intermediate_size();
        let up_proj = linear(n_embd, intermediate, vb.pp("up_proj"))?;
        let down_proj = linear(intermediate, n_embd, vb.pp("down_proj"))?;
        Ok(Self { up_proj, down_proj })
    }
}

impl Module for Mlp {
    fn forward(&self, x: &Tensor) -> CandleResult<Tensor> {
        let h = self.up_proj.forward(x)?;
        let h = h.gelu_erf()?;
        self.down_proj.forward(&h)
    }
}

/// Transformer Block — 标准 Pre-Norm GPT 结构
///
/// ```text
/// x ──────────────────────────┐
/// │                           │ (residual)
/// ├─ RMSNorm → Attention ─────┘ → +
/// │                           │ (residual)
/// └─ RMSNorm → MLP ──────────┘ → +
/// ```
struct TransformerBlock {
    ln1: RmsNorm,
    attn: CausalSelfAttention,
    ln2: RmsNorm,
    mlp: Mlp,
}

impl TransformerBlock {
    fn new(config: &GptConfig, vb: VarBuilder) -> CandleResult<Self> {
        Ok(Self {
            ln1: RmsNorm::new(config.n_embd, config.rms_norm_eps, vb.pp("ln1"))?,
            attn: CausalSelfAttention::new(config, vb.pp("attn"))?,
            ln2: RmsNorm::new(config.n_embd, config.rms_norm_eps, vb.pp("ln2"))?,
            mlp: Mlp::new(config, vb.pp("mlp"))?,
        })
    }

    /// x: (B, T, C) → (B, T, C)
    fn forward(&self, x: &Tensor, mask: &Tensor) -> CandleResult<Tensor> {
        // Attention with residual
        let h = self.ln1.forward(x)?;
        let h = self.attn.forward(&h, mask)?;
        let x = (x + h)?;
        // MLP with residual
        let h = self.ln2.forward(&x)?;
        let h = self.mlp.forward(&h)?;
        let x = (x + h)?;
        Ok(x)
    }
}

// ============================================================================
// T68.5 — MicroGpt: forward() + DeviceSelector
// ============================================================================

/// MicroGPT — 数据库原生 GPT 模型
///
/// 完整的 Decoder-only Transformer，支持推理和训练。
/// 通过 `DeviceSelector` 自动选择 CPU/CUDA/Metal。
///
/// # 使用示例
///
/// ```ignore
/// let config = GptConfig::tiny();
/// let device = Device::Cpu;
/// let gpt = MicroGpt::new(&config, &device)?;
///
/// // 推理
/// let tokens = vec![72u32, 101, 108, 108, 111]; // "Hello"
/// let gen_config = GenerateConfig { max_new_tokens: 20, temperature: 0.8, top_k: 40 };
/// let generated = gpt.generate(&tokens, &gen_config)?;
///
/// // 训练
/// let trainer = GptTrainer::new(&gpt, TrainConfig::default())?;
/// ```
pub struct MicroGpt {
    /// Token Embedding: vocab_size → n_embd
    wte: Embedding,
    /// Positional Embedding: block_size → n_embd
    wpe: Embedding,
    /// Transformer Blocks
    blocks: Vec<TransformerBlock>,
    /// Final RMSNorm
    ln_f: RmsNorm,
    /// Language Model Head: n_embd → vocab_size
    lm_head: Linear,
    /// 模型配置
    config: GptConfig,
    /// 设备
    device: Device,
    /// 可训练参数 (VarMap, 用于 optimizer)
    varmap: VarMap,
}

impl MicroGpt {
    /// 创建新的 MicroGPT 模型（随机初始化）
    ///
    /// # Arguments
    /// * `config` - 模型配置
    /// * `device` - 计算设备 (CPU/CUDA/Metal)
    pub fn new(config: &GptConfig, device: &Device) -> CandleResult<Self> {
        assert_eq!(
            config.n_embd % config.n_head,
            0,
            "n_embd ({}) must be divisible by n_head ({})",
            config.n_embd,
            config.n_head
        );

        let varmap = VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, device);

        // Token + Position Embedding
        let wte = embedding(config.vocab_size, config.n_embd, vb.pp("wte"))?;
        let wpe = embedding(config.block_size, config.n_embd, vb.pp("wpe"))?;

        // Transformer Blocks
        let mut blocks = Vec::with_capacity(config.n_layer);
        for i in 0..config.n_layer {
            blocks.push(TransformerBlock::new(config, vb.pp(format!("block.{i}")))?);
        }

        // Final normalization + LM Head
        let ln_f = RmsNorm::new(config.n_embd, config.rms_norm_eps, vb.pp("ln_f"))?;
        let lm_head = linear_no_bias(config.n_embd, config.vocab_size, vb.pp("lm_head"))?;

        Ok(Self {
            wte,
            wpe,
            blocks,
            ln_f,
            lm_head,
            config: config.clone(),
            device: device.clone(),
            varmap,
        })
    }

    /// 使用 DeviceSelector 自动选择最优设备并创建模型
    pub fn with_device_selector(
        config: &GptConfig,
        selector: &DeviceSelector,
    ) -> CandleResult<Self> {
        let (device, is_gpu) = selector.candle_device_with_fallback("microgpt_init");
        eprintln!(
            "[MicroGPT] initializing on {} (~{}K params, {}L/{}H/{}E)",
            if is_gpu { "GPU" } else { "CPU" },
            config.estimate_params() / 1000,
            config.n_layer,
            config.n_head,
            config.n_embd,
        );
        Self::new(config, &device)
    }

    /// 前向传播 — 输入 token IDs → 输出 logits
    ///
    /// # Arguments
    /// * `token_ids` - 形状 (B, T)，值域 `[0, vocab_size)`
    ///
    /// # Returns
    /// logits 张量，形状 (B, T, vocab_size)
    pub fn forward(&self, token_ids: &Tensor) -> CandleResult<Tensor> {
        let (_b, t) = token_ids.dims2()?;
        assert!(
            t <= self.config.block_size,
            "Sequence length {} exceeds block_size {}",
            t,
            self.config.block_size
        );

        // Position IDs: [0, 1, 2, ..., T-1]
        let pos_ids = Tensor::arange(0u32, t as u32, &self.device)?;

        // Token + Position Embedding
        let tok_emb = self.wte.forward(token_ids)?; // (B, T, C)
        let pos_emb = self.wpe.forward(&pos_ids)?; // (T, C)
        let mut x = tok_emb.broadcast_add(&pos_emb)?; // (B, T, C)

        // 因果注意力掩码
        let mask = causal_mask(t, &self.device)?;

        // N × Transformer Block
        for block in &self.blocks {
            x = block.forward(&x, &mask)?;
        }

        // Final RMSNorm + LM Head → logits
        let x = self.ln_f.forward(&x)?;
        self.lm_head.forward(&x)
    }

    /// 获取模型配置
    pub fn config(&self) -> &GptConfig {
        &self.config
    }

    /// 获取模型设备
    pub fn device(&self) -> &Device {
        &self.device
    }

    /// 获取 VarMap（用于外部 optimizer 管理）
    pub fn varmap(&self) -> &VarMap {
        &self.varmap
    }

    /// 估算模型参数量
    pub fn num_params(&self) -> usize {
        self.varmap
            .all_vars()
            .iter()
            .map(|v| v.elem_count())
            .sum()
    }
}

// ============================================================================
// T68.6 — train_step(): 损失计算 + 反向传播 + AdamW
// ============================================================================

/// 训练结果
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainStepResult {
    /// 本步损失值
    pub loss: f64,
    /// 训练步数
    pub step: usize,
    /// 每秒处理的 token 数
    pub tokens_per_sec: Option<f64>,
}

/// 训练配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainConfig {
    /// 学习率
    pub learning_rate: f64,
    /// AdamW 权重衰减
    pub weight_decay: f64,
    /// AdamW beta1
    pub beta1: f64,
    /// AdamW beta2
    pub beta2: f64,
    /// 梯度裁剪（最大范数）— 预留字段，当前未使用
    #[allow(dead_code)]
    pub max_grad_norm: Option<f64>,
}

impl Default for TrainConfig {
    fn default() -> Self {
        Self {
            learning_rate: 3e-4,
            weight_decay: 0.1,
            beta1: 0.9,
            beta2: 0.999,
            max_grad_norm: Some(1.0),
        }
    }
}

/// GPT 训练器 — 封装 optimizer 和训练循环
pub struct GptTrainer {
    /// AdamW 优化器
    optimizer: candle_nn::AdamW,
    /// 训练配置
    #[allow(dead_code)]
    config: TrainConfig,
    /// 当前步数
    step: usize,
}

impl GptTrainer {
    /// 创建训练器
    pub fn new(model: &MicroGpt, config: TrainConfig) -> CandleResult<Self> {
        let params = candle_nn::ParamsAdamW {
            lr: config.learning_rate,
            beta1: config.beta1,
            beta2: config.beta2,
            weight_decay: config.weight_decay,
            eps: 1e-8,
        };
        let optimizer = candle_nn::AdamW::new(model.varmap.all_vars(), params)?;
        Ok(Self {
            optimizer,
            config,
            step: 0,
        })
    }

    /// 执行单步训练
    ///
    /// # Arguments
    /// * `model` - MicroGPT 模型
    /// * `input_ids` - 输入 token IDs，形状 (B, T)
    /// * `target_ids` - 目标 token IDs，形状 (B, T)（通常为 input_ids 左移1位）
    ///
    /// # Returns
    /// 本步训练结果（损失值等）
    pub fn train_step(
        &mut self,
        model: &MicroGpt,
        input_ids: &Tensor,
        target_ids: &Tensor,
    ) -> CandleResult<TrainStepResult> {
        let start = std::time::Instant::now();
        let (b, t) = input_ids.dims2()?;

        // 前向传播 → logits (B, T, V)
        let logits = model.forward(input_ids)?;

        // Cross-entropy loss
        // 展平: logits (B*T, V), targets (B*T)
        let logits_flat = logits.reshape((b * t, model.config.vocab_size))?;
        let targets_flat = target_ids.reshape(b * t)?;
        let loss = candle_nn::loss::cross_entropy(&logits_flat, &targets_flat)?;

        // 反向传播
        let grads = loss.backward()?;
        // AdamW optimizer step
        self.optimizer.step(&grads)?;

        self.step += 1;
        let elapsed = start.elapsed().as_secs_f64();
        let tokens_per_sec = (b * t) as f64 / elapsed;

        Ok(TrainStepResult {
            loss: loss.to_scalar::<f32>()? as f64,
            step: self.step,
            tokens_per_sec: Some(tokens_per_sec),
        })
    }

    /// 从 token 序列生成训练对（input, target）
    ///
    /// 将序列切分为 (tokens\[..T\], tokens\[1..T+1\]) 的输入/目标对。
    pub fn prepare_batch(
        tokens: &[u32],
        block_size: usize,
        batch_size: usize,
        device: &Device,
    ) -> CandleResult<(Tensor, Tensor)> {
        let n = tokens.len();
        if n < block_size + 1 {
            return Err(candle_core::Error::Msg(format!(
                "Token sequence length ({}) must be >= block_size + 1 ({})",
                n,
                block_size + 1
            )));
        }

        let mut input_data = Vec::with_capacity(batch_size * block_size);
        let mut target_data = Vec::with_capacity(batch_size * block_size);

        let mut rng = rand::thread_rng();
        for _ in 0..batch_size {
            let start = rng.gen_range(0..n - block_size);
            input_data.extend_from_slice(&tokens[start..start + block_size]);
            target_data.extend_from_slice(&tokens[start + 1..start + 1 + block_size]);
        }

        let input_ids = Tensor::from_vec(input_data, (batch_size, block_size), device)?;
        let target_ids = Tensor::from_vec(target_data, (batch_size, block_size), device)?;
        Ok((input_ids, target_ids))
    }

    /// 获取当前步数
    pub fn step(&self) -> usize {
        self.step
    }

    /// 获取训练配置
    pub fn train_config(&self) -> &TrainConfig {
        &self.config
    }
}

// ============================================================================
// T68.7 — generate(): 自回归温度采样推理
// ============================================================================

/// 生成配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateConfig {
    /// 生成的最大 token 数
    pub max_new_tokens: usize,
    /// 温度（越高越随机，1.0=标准，0.0=贪心）
    pub temperature: f64,
    /// Top-K 采样（0=关闭）
    pub top_k: usize,
}

impl Default for GenerateConfig {
    fn default() -> Self {
        Self {
            max_new_tokens: 100,
            temperature: 0.8,
            top_k: 40,
        }
    }
}

impl MicroGpt {
    /// 自回归生成文本 token
    ///
    /// 从给定的 prompt tokens 开始，逐步生成新 token，每次将新 token
    /// 追加到序列中并再次前向传播，直到达到 max_new_tokens。
    ///
    /// # Arguments
    /// * `prompt_tokens` - 初始 prompt 的 token IDs
    /// * `gen_config` - 生成配置
    ///
    /// # Returns
    /// 完整序列 (prompt + generated)
    pub fn generate(
        &self,
        prompt_tokens: &[u32],
        gen_config: &GenerateConfig,
    ) -> CandleResult<Vec<u32>> {
        let mut tokens = prompt_tokens.to_vec();

        for _ in 0..gen_config.max_new_tokens {
            // 裁剪到 block_size 窗口
            let start = if tokens.len() > self.config.block_size {
                tokens.len() - self.config.block_size
            } else {
                0
            };
            let context = &tokens[start..];

            // 前向传播: (1, T) → (1, T, V)
            let input =
                Tensor::from_vec(context.to_vec(), (1, context.len()), &self.device)?;
            let logits = self.forward(&input)?;

            // 取最后一个位置的 logits: (V,)
            let logits = logits.i((0, context.len() - 1))?;

            // 采样下一个 token
            let next_token =
                sample_token(&logits, gen_config.temperature, gen_config.top_k)?;
            tokens.push(next_token);
        }

        Ok(tokens)
    }
}

/// Top-K 温度采样
///
/// 1. 应用温度缩放: logits / temperature
/// 2. Top-K 过滤: 只保留概率最高的 K 个 token
/// 3. Softmax → 概率分布
/// 4. 按概率采样
fn sample_token(logits: &Tensor, temperature: f64, top_k: usize) -> CandleResult<u32> {
    if temperature <= 0.0 || temperature < 1e-10 {
        // 贪心解码
        let idx = logits.argmax(0)?.to_scalar::<u32>()?;
        return Ok(idx);
    }

    let logits = (logits / temperature)?;
    let logits_vec: Vec<f32> = logits.to_vec1()?;
    let vocab_size = logits_vec.len();

    // Top-K 过滤
    let (filtered_logits, filtered_indices) = if top_k > 0 && top_k < vocab_size {
        let mut indexed: Vec<(usize, f32)> =
            logits_vec.iter().copied().enumerate().collect();
        indexed
            .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        let top: Vec<(usize, f32)> = indexed.into_iter().take(top_k).collect();
        let indices: Vec<usize> = top.iter().map(|(i, _)| *i).collect();
        let vals: Vec<f32> = top.iter().map(|(_, v)| *v).collect();
        (vals, indices)
    } else {
        let indices: Vec<usize> = (0..vocab_size).collect();
        (logits_vec, indices)
    };

    // Softmax (CPU, numerically stable)
    let max_val = filtered_logits
        .iter()
        .cloned()
        .fold(f32::NEG_INFINITY, f32::max);
    let exp_vals: Vec<f32> = filtered_logits
        .iter()
        .map(|&v| (v - max_val).exp())
        .collect();
    let sum: f32 = exp_vals.iter().sum();
    let probs: Vec<f32> = exp_vals.iter().map(|&v| v / sum).collect();

    // 按概率采样
    let r: f32 = rand::thread_rng().gen::<f32>();
    let mut cumsum = 0.0f32;
    for (i, &p) in probs.iter().enumerate() {
        cumsum += p;
        if r <= cumsum {
            return Ok(filtered_indices[i] as u32);
        }
    }

    // 兜底返回最后一个
    Ok(*filtered_indices.last().unwrap_or(&0) as u32)
}

// ============================================================================
// T68.8 — TokenizerManager: 字符级 + BPE
// ============================================================================

/// Tokenizer 类型
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum TokenizerType {
    /// 字符级 tokenizer（每个字节=一个 token，vocab_size=256）
    CharLevel,
    /// BPE tokenizer（通过 HuggingFace tokenizers crate）
    Bpe {
        /// 模型 ID 或本地路径
        model_id: String,
    },
}

/// Tokenizer 管理器
///
/// 统一接口支持字符级和 BPE 两种模式。
/// 字符级模式无需外部模型文件，适合快速原型验证；
/// BPE 模式通过 `tokenizers` crate 加载 HuggingFace 模型。
pub struct TokenizerManager {
    tokenizer_type: TokenizerType,
    bpe_tokenizer: Option<tokenizers::Tokenizer>,
}

impl TokenizerManager {
    /// 创建字符级 tokenizer
    ///
    /// 每个字节直接映射为 token ID (0-255)，vocab_size = 256。
    /// 简单高效，适合训练小模型。
    pub fn char_level() -> Self {
        Self {
            tokenizer_type: TokenizerType::CharLevel,
            bpe_tokenizer: None,
        }
    }

    /// 从本地 tokenizer.json 文件加载 BPE tokenizer
    pub fn from_file(path: &str) -> Result<Self, String> {
        let tokenizer = tokenizers::Tokenizer::from_file(path)
            .map_err(|e| format!("Failed to load tokenizer from '{}': {}", path, e))?;

        Ok(Self {
            tokenizer_type: TokenizerType::Bpe {
                model_id: path.to_string(),
            },
            bpe_tokenizer: Some(tokenizer),
        })
    }

    /// 从已加载的 tokenizers::Tokenizer 实例创建
    pub fn from_tokenizer(tokenizer: tokenizers::Tokenizer, name: &str) -> Self {
        Self {
            tokenizer_type: TokenizerType::Bpe {
                model_id: name.to_string(),
            },
            bpe_tokenizer: Some(tokenizer),
        }
    }

    /// 编码文本为 token IDs
    pub fn encode(&self, text: &str) -> Result<Vec<u32>, String> {
        match &self.tokenizer_type {
            TokenizerType::CharLevel => {
                Ok(text.as_bytes().iter().map(|&b| b as u32).collect())
            }
            TokenizerType::Bpe { .. } => {
                let tokenizer = self
                    .bpe_tokenizer
                    .as_ref()
                    .ok_or("BPE tokenizer not loaded")?;
                let encoding = tokenizer
                    .encode(text, false)
                    .map_err(|e| format!("Encode error: {}", e))?;
                Ok(encoding.get_ids().to_vec())
            }
        }
    }

    /// 解码 token IDs 为文本
    pub fn decode(&self, token_ids: &[u32]) -> Result<String, String> {
        match &self.tokenizer_type {
            TokenizerType::CharLevel => {
                let bytes: Vec<u8> = token_ids.iter().map(|&id| id as u8).collect();
                Ok(String::from_utf8_lossy(&bytes).to_string())
            }
            TokenizerType::Bpe { .. } => {
                let tokenizer = self
                    .bpe_tokenizer
                    .as_ref()
                    .ok_or("BPE tokenizer not loaded")?;
                tokenizer
                    .decode(token_ids, true)
                    .map_err(|e| format!("Decode error: {}", e))
            }
        }
    }

    /// 获取词表大小
    pub fn vocab_size(&self) -> usize {
        match &self.tokenizer_type {
            TokenizerType::CharLevel => 256,
            TokenizerType::Bpe { .. } => self
                .bpe_tokenizer
                .as_ref()
                .map(|t: &tokenizers::Tokenizer| t.get_vocab_size(true))
                .unwrap_or(0),
        }
    }

    /// 获取 tokenizer 类型
    pub fn tokenizer_type(&self) -> &TokenizerType {
        &self.tokenizer_type
    }
}

impl std::fmt::Debug for TokenizerManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TokenizerManager")
            .field("type", &self.tokenizer_type)
            .field("vocab_size", &self.vocab_size())
            .finish()
    }
}

// ============================================================================
// 便捷 API — 端到端训练/推理
// ============================================================================

/// 端到端训练结果
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingReport {
    /// 初始损失
    pub initial_loss: f64,
    /// 最终损失
    pub final_loss: f64,
    /// 总训练步数
    pub total_steps: usize,
    /// 损失历史
    pub loss_history: Vec<f64>,
    /// 总训练时间（秒）
    pub elapsed_secs: f64,
    /// 平均 tokens/sec
    pub avg_tokens_per_sec: f64,
}

impl MicroGpt {
    /// 端到端训练：从文本数据训练模型
    ///
    /// # Arguments
    /// * `text` - 训练文本
    /// * `tokenizer` - tokenizer 管理器
    /// * `train_config` - 训练超参数
    /// * `num_steps` - 训练步数
    /// * `batch_size` - 批大小
    ///
    /// # Returns
    /// 训练报告
    pub fn train_on_text(
        &self,
        text: &str,
        tokenizer: &TokenizerManager,
        train_config: TrainConfig,
        num_steps: usize,
        batch_size: usize,
    ) -> CandleResult<TrainingReport> {
        let tokens = tokenizer
            .encode(text)
            .map_err(candle_core::Error::Msg)?;

        if tokens.len() < self.config.block_size + 1 {
            return Err(candle_core::Error::Msg(format!(
                "Training text too short: {} tokens, need at least {}",
                tokens.len(),
                self.config.block_size + 1
            )));
        }

        let mut trainer = GptTrainer::new(self, train_config)?;
        let mut loss_history = Vec::with_capacity(num_steps);
        let start = std::time::Instant::now();
        let mut total_tokens = 0usize;

        for _step in 0..num_steps {
            let (input_ids, target_ids) = GptTrainer::prepare_batch(
                &tokens,
                self.config.block_size,
                batch_size,
                &self.device,
            )?;
            let result = trainer.train_step(self, &input_ids, &target_ids)?;
            loss_history.push(result.loss);
            total_tokens += batch_size * self.config.block_size;
        }

        let elapsed = start.elapsed().as_secs_f64();
        Ok(TrainingReport {
            initial_loss: *loss_history.first().unwrap_or(&0.0),
            final_loss: *loss_history.last().unwrap_or(&0.0),
            total_steps: num_steps,
            loss_history,
            elapsed_secs: elapsed,
            avg_tokens_per_sec: total_tokens as f64 / elapsed,
        })
    }

    /// 端到端生成：从文本 prompt 生成文本
    ///
    /// # Arguments
    /// * `prompt` - 文本 prompt
    /// * `tokenizer` - tokenizer 管理器
    /// * `gen_config` - 生成配置
    ///
    /// # Returns
    /// 生成的文本
    pub fn generate_text(
        &self,
        prompt: &str,
        tokenizer: &TokenizerManager,
        gen_config: &GenerateConfig,
    ) -> CandleResult<String> {
        let prompt_tokens = tokenizer
            .encode(prompt)
            .map_err(candle_core::Error::Msg)?;

        let output_tokens = self.generate(&prompt_tokens, gen_config)?;

        let text = tokenizer
            .decode(&output_tokens)
            .map_err(candle_core::Error::Msg)?;

        Ok(text)
    }
}

// ============================================================================
// 测试
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    // --- T68.1 Tests: Config ---

    #[test]
    fn test_gpt_config_default() {
        let config = GptConfig::default();
        assert_eq!(config.vocab_size, 256);
        assert_eq!(config.n_embd, 64);
        assert_eq!(config.n_head, 4);
        assert_eq!(config.head_dim(), 16);
        assert_eq!(config.intermediate_size(), 256);
        assert!(config.estimate_params() > 0);
    }

    #[test]
    fn test_gpt_config_tiny() {
        let config = GptConfig::tiny();
        assert_eq!(config.n_embd % config.n_head, 0);
        println!("Tiny config params: ~{}K", config.estimate_params() / 1000);
    }

    // --- T68.2 Tests: RmsNorm + causal_mask ---

    #[test]
    fn test_causal_mask() {
        let mask = causal_mask(4, &Device::Cpu).unwrap();
        assert_eq!(mask.dims(), &[1, 1, 4, 4]);
        let vals: Vec<Vec<f32>> = (0..4)
            .map(|i| mask.i((0, 0, i)).unwrap().to_vec1::<f32>().unwrap())
            .collect();
        // 对角线和下方应该为 0，上方应该为 -inf
        assert_eq!(vals[0][0], 0.0);
        assert!(vals[0][1].is_infinite() && vals[0][1] < 0.0);
        assert_eq!(vals[1][0], 0.0);
        assert_eq!(vals[1][1], 0.0);
        assert!(vals[1][2].is_infinite() && vals[1][2] < 0.0);
        assert_eq!(vals[3][3], 0.0);
    }

    // --- T68.3~T68.5 Tests: Forward Pass ---

    #[test]
    fn test_micro_gpt_forward() {
        let config = GptConfig::tiny();
        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();

        // 输入: batch=2, seq_len=8
        let input = Tensor::from_vec(vec![0u32; 16], (2, 8), &device).unwrap();

        let logits = model.forward(&input).unwrap();
        assert_eq!(logits.dims(), &[2, 8, config.vocab_size]);
        println!("Forward pass OK, logits shape: {:?}", logits.dims());
    }

    #[test]
    fn test_micro_gpt_num_params() {
        let config = GptConfig::tiny();
        let model = MicroGpt::new(&config, &Device::Cpu).unwrap();
        let actual = model.num_params();
        let estimated = config.estimate_params();
        println!("Actual params: {}, Estimated: {}", actual, estimated);
        assert!(actual > 0);
    }

    #[test]
    fn test_micro_gpt_forward_different_seq_lens() {
        let config = GptConfig::tiny();
        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();

        for seq_len in [1, 4, 16, config.block_size] {
            let input =
                Tensor::from_vec(vec![42u32; seq_len], (1, seq_len), &device).unwrap();
            let logits = model.forward(&input).unwrap();
            assert_eq!(logits.dims(), &[1, seq_len, config.vocab_size]);
        }
    }

    // --- T68.6 Tests: Training ---

    #[test]
    fn test_train_step_loss_decreases() {
        let config = GptConfig::tiny();
        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();
        let train_config = TrainConfig {
            learning_rate: 1e-3,
            ..Default::default()
        };
        let mut trainer = GptTrainer::new(&model, train_config).unwrap();

        // 简单训练数据: 重复序列 [0,1,2,...,63]
        let tokens: Vec<u32> = (0..256).map(|i| (i % 64) as u32).collect();
        let block_size = config.block_size;

        let mut losses = Vec::new();
        for _ in 0..10 {
            let (input, target) =
                GptTrainer::prepare_batch(&tokens, block_size, 2, &device).unwrap();
            let result = trainer.train_step(&model, &input, &target).unwrap();
            losses.push(result.loss);
        }

        // 损失应该有下降趋势
        println!("Training losses: {:?}", losses);
        let first_3_avg: f64 = losses[..3].iter().sum::<f64>() / 3.0;
        let last_3_avg: f64 = losses[7..].iter().sum::<f64>() / 3.0;
        assert!(
            last_3_avg < first_3_avg,
            "Loss should decrease: first_3_avg={:.4}, last_3_avg={:.4}",
            first_3_avg, last_3_avg
        );
    }

    #[test]
    fn test_prepare_batch() {
        let tokens: Vec<u32> = (0..200).collect();
        let (input, target) =
            GptTrainer::prepare_batch(&tokens, 32, 4, &Device::Cpu).unwrap();
        assert_eq!(input.dims(), &[4, 32]);
        assert_eq!(target.dims(), &[4, 32]);
    }

    // --- T68.7 Tests: Generation ---

    #[test]
    fn test_generate_greedy() {
        let config = GptConfig::tiny();
        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();

        let prompt = vec![72u32, 101, 108, 108, 111]; // "Hello"
        let gen_config = GenerateConfig {
            max_new_tokens: 10,
            temperature: 0.0, // 贪心
            top_k: 0,
        };
        let output = model.generate(&prompt, &gen_config).unwrap();
        assert_eq!(output.len(), 15); // 5 prompt + 10 generated
        // 贪心解码应该是确定性的
        let output2 = model.generate(&prompt, &gen_config).unwrap();
        assert_eq!(output, output2, "Greedy generation should be deterministic");
    }

    #[test]
    fn test_generate_with_temperature() {
        let config = GptConfig::tiny();
        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();

        let prompt = vec![72u32, 101, 108];
        let gen_config = GenerateConfig {
            max_new_tokens: 20,
            temperature: 0.8,
            top_k: 40,
        };
        let output = model.generate(&prompt, &gen_config).unwrap();
        assert_eq!(output.len(), 23);
        // 前 3 个应该是 prompt
        assert_eq!(&output[..3], &prompt);
    }

    // --- T68.8 Tests: Tokenizer ---

    #[test]
    fn test_char_level_tokenizer() {
        let tok = TokenizerManager::char_level();
        assert_eq!(tok.vocab_size(), 256);

        let text = "Hello, World!";
        let tokens = tok.encode(text).unwrap();
        assert_eq!(tokens.len(), text.len());
        assert_eq!(tokens[0], b'H' as u32);
        assert_eq!(tokens[1], b'e' as u32);

        let decoded = tok.decode(&tokens).unwrap();
        assert_eq!(decoded, text);
    }

    #[test]
    fn test_char_level_roundtrip() {
        let tok = TokenizerManager::char_level();
        let texts = vec!["Simple ASCII text", "Numbers: 12345", "Special: !@#$%", ""];
        for text in texts {
            let tokens = tok.encode(text).unwrap();
            let decoded = tok.decode(&tokens).unwrap();
            assert_eq!(decoded, text, "Roundtrip failed for: {:?}", text);
        }
    }

    // --- 端到端测试 ---

    #[test]
    fn test_end_to_end_train_and_generate() {
        let config = GptConfig::tiny();
        let device = Device::Cpu;
        let model = MicroGpt::new(&config, &device).unwrap();
        let tokenizer = TokenizerManager::char_level();

        // 训练文本（重复模式，便于学习）
        let train_text = "abcdefgh".repeat(50);

        let report = model
            .train_on_text(
                &train_text,
                &tokenizer,
                TrainConfig {
                    learning_rate: 1e-3,
                    ..Default::default()
                },
                20, // steps
                2,  // batch_size
            )
            .unwrap();

        println!(
            "Training report: loss {:.4} → {:.4}, {:.0} tok/s",
            report.initial_loss, report.final_loss, report.avg_tokens_per_sec
        );
        assert!(
            report.final_loss < report.initial_loss,
            "Loss should decrease"
        );

        // 生成文本
        let gen_config = GenerateConfig {
            max_new_tokens: 30,
            temperature: 0.5,
            top_k: 10,
        };
        let output = model
            .generate_text("abc", &tokenizer, &gen_config)
            .unwrap();
        println!("Generated: {:?}", output);
        assert!(output.starts_with("abc"));
        assert!(output.len() > 3);
    }

    #[test]
    fn test_device_selector_integration() {
        let config = GptConfig::tiny();
        let selector = DeviceSelector::auto();
        let model = MicroGpt::with_device_selector(&config, &selector).unwrap();

        let input =
            Tensor::from_vec(vec![0u32; 8], (1, 8), model.device()).unwrap();
        let logits = model.forward(&input).unwrap();
        assert_eq!(logits.dims(), &[1, 8, config.vocab_size]);
    }
}
