import torch
import torch.nn as nn
import lightning as L

from gpt.tokenizer import token_ids_to_text
from gpt.layer import TransformerBlock, LayerNorm


class GPTModel(L.LightningModule):
    """
    GPT 语言模型，基于 Transformer 解码器架构。

    模型结构：
        - Token Embedding：将 token ID 映射为 d_model 维向量
        - Position Embedding：为每个位置添加可学习的位置编码
        - Dropout：对嵌入向量进行正则化
        - Transformer 块堆叠：多个 TransformerBlock
        - 最终层归一化
        - 输出线性层（LM Head）：将 d_model 映射回词汇表大小

    Args:
        cfg (dict): 配置字典，必须包含以下键：
            - vocab_size: 词汇表大小
            - d_model: 嵌入维度
            - context_length: 最大上下文长度
            - drop_rate: Dropout 概率
            - n_layers: Transformer 块的数量
    """

    def __init__(self, cfg):
        super().__init__()

        # Token 嵌入层
        self.token_embedding = nn.Embedding(cfg["vocab_size"], cfg["d_model"])

        # 位置嵌入层
        self.position_embedding = nn.Embedding(cfg["context_length"], cfg["d_model"])

        # 嵌入后的 Dropout
        self.embedding_dropout = nn.Dropout(cfg["drop_rate"])

        # Transformer 块堆叠
        self.transformer_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
        )

        # 最终层归一化
        self.final_layer_norm = LayerNorm(cfg["d_model"])

        # 输出层（语言模型头），无偏置
        self.lm_head = nn.Linear(cfg["d_model"], cfg["vocab_size"], bias=False)

        # 用于记录训练和验证过程中的损失（用于后续绘图等）
        self.train_step_losses = []  # 每个 step 的损失
        self.train_epoch_losses = []  # 每个 epoch 的平均损失
        self.val_step_losses = []  # 每个 step 的验证损失
        self.val_epoch_losses = []  # 每个 epoch 的平均验证损失

    def forward(self, input_ids):
        """
        前向传播，计算下一个 token 的 logits。

        Args:
            input_ids (torch.Tensor): 输入 token ID 序列，形状为 (batch_size, seq_len)

        Returns:
            torch.Tensor: 预测 logits，形状为 (batch_size, seq_len, vocab_size)
        """
        batch_size, seq_len = input_ids.shape
        device = input_ids.device

        # Token 嵌入
        token_embeds = self.token_embedding(input_ids)  # (batch, seq_len, d_model)

        # 位置嵌入（生成位置索引 [0, 1, ..., seq_len-1]）
        position_ids = torch.arange(seq_len, device=device)
        position_embeds = self.position_embedding(position_ids)  # (seq_len, d_model)

        # 将 token 嵌入和位置嵌入相加，得到最终的输入表示
        hidden_states = token_embeds + position_embeds  # (batch, seq_len, d_model)

        # 应用 Dropout
        hidden_states = self.embedding_dropout(hidden_states)

        # 通过所有 Transformer 块
        hidden_states = self.transformer_blocks(hidden_states)

        # 最终层归一化
        hidden_states = self.final_layer_norm(hidden_states)

        # 通过语言模型头得到 logits
        logits = self.lm_head(hidden_states)  # (batch, seq_len, vocab_size)

        return logits

    def training_step(self, batch, batch_idx):
        """
        单个训练步骤。计算损失并记录。

        Args:
            batch: 由 DataLoader 返回的批次，包含 (input_ids, target_ids)
            batch_idx: 批次索引

        Returns:
            torch.Tensor: 损失值
        """
        input_ids, target_ids = batch

        logits = self(input_ids)  # (batch, seq_len, vocab_size)

        # 计算交叉熵损失（将 logits 和 targets 展平）
        loss = nn.functional.cross_entropy(
            logits.view(-1, logits.size(-1)), target_ids.view(-1)
        )

        # 记录当前 step 的损失（detach 避免计算图梯度累积）
        self.train_step_losses.append(loss.detach())

        # 记录到 Lightning 日志（可在 TensorBoard 中查看）
        self.log("train_step_loss", loss.detach())

        return loss

    def on_train_epoch_end(self):
        """每个训练 epoch 结束时计算平均损失并记录。"""
        if self.train_step_losses:
            avg_loss = torch.stack(self.train_step_losses).mean()
            self.print(
                f"***** 【Epoch {self.current_epoch}】 Train Avg Loss: {avg_loss:.4f} *****"
            )
            self.train_epoch_losses.append(
                {"epoch": self.current_epoch, "loss": avg_loss.item()}
            )
            # 清空 step 列表，为下一 epoch 做准备
            self.train_step_losses.clear()

    def validation_step(self, batch, batch_idx):
        """单个验证步骤。"""
        input_ids, target_ids = batch
        logits = self(input_ids)
        loss = nn.functional.cross_entropy(
            logits.view(-1, logits.size(-1)), target_ids.view(-1)
        )
        self.val_step_losses.append(loss.detach())
        self.log("val_step_loss", loss.detach(), on_step=True, on_epoch=False)
        return loss

    def on_validation_epoch_end(self):
        """每个验证 epoch 结束时计算平均损失并记录。"""
        if self.val_step_losses:
            avg_loss = torch.stack(self.val_step_losses).mean()
            self.print(
                f"***** 【Epoch {self.current_epoch}】 Val Avg Loss: {avg_loss:.4f} *****"
            )
            self.val_epoch_losses.append(
                {"epoch": self.current_epoch, "loss": avg_loss.item()}
            )
            # 清空 step 列表，为下一 epoch 做准备
            self.val_step_losses.clear()

    def configure_optimizers(self):
        """配置优化器（AdamW）。"""
        optimizer = torch.optim.AdamW(self.parameters(), lr=0.0004, weight_decay=0.01)
        return optimizer

    def clear_loss_history(self):
        """清空所有缓存的损失列表，释放内存。"""
        self.train_step_losses.clear()
        self.train_epoch_losses.clear()
        self.val_step_losses.clear()
        self.val_epoch_losses.clear()


def greedy_generate(model, start_ids, max_new_tokens, context_size, tokenizer):
    """
    使用贪心解码策略生成文本。

    Args:
        model: 语言模型（应处于 eval 模式）
        start_ids (torch.Tensor): 输入的提示 token ID 序列，形状为 (batch_size, seq_len)
        max_new_tokens (int): 要生成的最大新 token 数量
        context_size (int): 模型支持的最大上下文长度
        tokenizer: 用于将 token IDs 解码为文本的分词器

    Returns:
        str: 生成的完整文本（包含原始提示）
    """
    model.eval()
    generated_ids = start_ids.to(model.device)

    for _ in range(max_new_tokens):
        # 如果当前序列长度超过上下文窗口，仅保留最后 context_size 个 token 作为输入
        input_context = generated_ids[:, -context_size:]

        with torch.no_grad():
            logits = model(input_context)  # (batch, seq_len, vocab_size)
            next_token_logits = logits[:, -1, :]  # (batch, vocab_size)

        # 选择概率最高的 token ID（贪心）
        next_token_id = torch.argmax(
            next_token_logits, dim=-1, keepdim=True
        )  # (batch, 1)

        # 拼接到已生成序列
        generated_ids = torch.cat((generated_ids, next_token_id), dim=1)

    # 将 token IDs 转换为文本（假设 batch_size = 1）
    return token_ids_to_text(generated_ids, tokenizer)
