#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LLM对话AI训练代码 - 改进版
基于Transformer架构，参考MiniMind/nanoGPT等开源项目优化
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import math
import json
import os
from typing import List, Tuple, Dict, Any
import time
import threading
import queue
import re
from collections import defaultdict


# ==================== 1. 字符级分词器（对中文友好）====================
class CharTokenizer:
    """简单字符级分词器 - 中文按字符，英文按单词
    
    每个字符/单词作为一个token，对中文效果更好。
    """
    
    def __init__(self):
        self.special_tokens = {
            '<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
            '<START>': 4, '<END>': 5, '<SEP>': 6
        }
        self.char_to_id = self.special_tokens.copy()
        self.id_to_char = {v: k for k, v in self.char_to_id.items()}
        self.next_id = len(self.special_tokens)
    
    def build_vocab(self, texts: List[str], min_freq: int = 1) -> Dict[str, int]:
        """从文本构建词汇表"""
        char_counts = defaultdict(int)
        
        for text in texts:
            # 中文按字符，英文按空格分割
            chars = self._split_chars(text)
            for char in chars:
                char_counts[char] += 1
        
        # 添加高频字符
        for char, count in char_counts.items():
            if count >= min_freq and char not in self.char_to_id:
                self.char_to_id[char] = self.next_id
                self.id_to_char[self.next_id] = char
                self.next_id += 1
        
        return self.char_to_id
    
    def _split_chars(self, text: str) -> List[str]:
        """分割文本为字符/词"""
        # 先尝试按空格分割英文
        parts = text.split(' ')
        result = []
        for part in parts:
            if part.isascii():
                result.append(part)
            else:
                # 中文字符逐个添加
                for char in part:
                    if char.strip():
                        result.append(char)
        return result
    
    def encode(self, text: str) -> List[int]:
        """将文本编码为ID序列"""
        if not text:
            return [self.special_tokens['<UNK>']]
        
        chars = self._split_chars(text)
        ids = []
        for char in chars:
            if char in self.char_to_id:
                ids.append(self.char_to_id[char])
            else:
                # 未知字符用 UNK token
                ids.append(self.special_tokens['<UNK>'])
        
        return ids
    
    def decode(self, ids: List[int]) -> str:
        """将ID序列解码为文本"""
        chars = []
        for tid in ids:
            if tid in self.id_to_char:
                char = self.id_to_char[tid]
                if char not in self.special_tokens:
                    chars.append(char)
        return ''.join(chars)


# 全局分词器实例
tokenizer_instance = CharTokenizer()


def build_vocabulary(conversations: List[Tuple[str, str]], min_freq: int = 1) -> dict:
    """构建词汇表"""
    all_texts = []
    for question, answer in conversations:
        all_texts.append(question)
        all_texts.append(answer)
    return tokenizer_instance.build_vocab(all_texts, min_freq)


# ==================== 2. 位置编码 ====================
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """预计算旋转位置编码的频率"""
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cos = torch.cos(freqs)
    freqs_sin = torch.sin(freqs)
    return freqs_cos, freqs_sin


def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    """调整频率张量形状以进行广播"""
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(shape)


def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor):
    """应用旋转位置编码 - 参考MiniMind实现"""
    # 将xq和xk的最后一个维度分成两半
    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
    
    # 调整freqs形状
    freqs_cos = reshape_for_broadcast(freqs_cos, xq_r)
    freqs_sin = reshape_for_broadcast(freqs_sin, xq_r)
    
    # 应用旋转
    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos
    
    # 重新组合
    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3)
    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3)
    
    return xq_out.type_as(xq), xk_out.type_as(xk)


# ==================== 3. RMSNorm ====================
class RMSNorm(nn.Module):
    """RMSNorm归一化层"""
    
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        rms = torch.sqrt(torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps)
        return self.weight * x / rms


# ==================== 4. 注意力机制 ====================
class Attention(nn.Module):
    """多头注意力 - 参考MiniMind优化实现"""
    
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        
        # QKV投影合并为一个矩阵（更高效）
        self.wq = nn.Linear(d_model, d_model, bias=False)
        self.wk = nn.Linear(d_model, d_model, bias=False)
        self.wv = nn.Linear(d_model, d_model, bias=False)
        self.wo = nn.Linear(d_model, d_model, bias=False)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, freqs_cos, freqs_sin, mask=None):
        batch_size, seq_len, _ = x.shape
        
        # QKV投影
        xq = self.wq(x)
        xk = self.wk(x)
        xv = self.wv(x)
        
        # 重塑为多头 [batch, seq, n_heads, d_head]
        xq = xq.view(batch_size, seq_len, self.n_heads, self.d_head)
        xk = xk.view(batch_size, seq_len, self.n_heads, self.d_head)
        xv = xv.view(batch_size, seq_len, self.n_heads, self.d_head)
        
        # 应用RoPE
        xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
        
        # 转置为 [batch, n_heads, seq, d_head]
        xq = xq.transpose(1, 2)
        xk = xk.transpose(1, 2)
        xv = xv.transpose(1, 2)
        
        # 计算注意力分数
        scores = torch.matmul(xq, xk.transpose(-2, -1)) / math.sqrt(self.d_head)
        
        # 应用掩码
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Softmax
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # 注意力输出
        output = torch.matmul(attn_weights, xv)
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        
        return self.wo(output)


# ==================== 5. 前馈网络 ====================
class FeedForward(nn.Module):
    """SwiGLU前馈网络 - 现代LLM标准实现"""
    
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        return self.w2(self.dropout(nn.functional.silu(self.w1(x)) * self.w3(x)))


# ==================== 6. MoE专家混合 ====================
class FeedForwardExpert(nn.Module):
    """MoE专家 - 单个前馈网络"""
    
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        return self.w2(self.dropout(nn.functional.silu(self.w1(x)) * self.w3(x)))


class MixtureOfExperts(nn.Module):
    """专家混合(MoE)层 - 高效并行实现"""
    
    def __init__(self, d_model: int, d_ff: int, num_experts: int = 4, top_k: int = 2, dropout: float = 0.1):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.d_model = d_model
        
        # 专家网络
        self.experts = nn.ModuleList([
            FeedForwardExpert(d_model, d_ff, dropout) for _ in range(num_experts)
        ])
        
        # 门控网络
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # [batch*seq, d_model]
        
        # 计算门控权重
        gate_logits = self.gate(x_flat)  # [batch*seq, num_experts]
        
        # 选择top-k专家
        top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)
        top_k_weights = torch.softmax(top_k_logits, dim=-1)  # [batch*seq, top_k]
        
        # 高效并行计算：为每个专家收集输入
        output = torch.zeros_like(x_flat)
        
        for expert_idx in range(self.num_experts):
            # 找到选择该专家的所有位置
            mask = (top_k_indices == expert_idx).any(dim=-1)  # [batch*seq]
            if not mask.any():
                continue
            
            # 获取选择该专家的输入
            expert_input = x_flat[mask]  # [num_selected, d_model]
            
            # 计算该专家的权重
            expert_positions = mask.nonzero(as_tuple=True)[0]
            expert_weights = []
            for pos in expert_positions:
                # 找到该位置在top_k中的权重
                pos_in_topk = (top_k_indices[pos] == expert_idx).nonzero(as_tuple=True)[0]
                if len(pos_in_topk) > 0:
                    expert_weights.append(top_k_weights[pos, pos_in_topk[0]])
                else:
                    expert_weights.append(torch.tensor(0.0, device=x.device))
            expert_weights = torch.stack(expert_weights).unsqueeze(-1)  # [num_selected, 1]
            
            # 前向传播
            expert_output = self.experts[expert_idx](expert_input) * expert_weights
            
            # 累加到输出
            output[mask] += expert_output
        
        return output.view(batch_size, seq_len, d_model)


# ==================== 7. Transformer块 ====================
class TransformerBlock(nn.Module):
    """Transformer块 - Pre-Norm + 可选MoE"""
    
    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1,
                 use_moe: bool = False, num_experts: int = 4, top_k: int = 2):
        super().__init__()
        
        self.attention = Attention(d_model, n_heads, dropout)
        
        # Pre-Norm
        self.attn_norm = RMSNorm(d_model)
        self.ff_norm = RMSNorm(d_model)
        
        # 前馈网络（可选MoE）
        if use_moe:
            self.feed_forward = MixtureOfExperts(d_model, d_ff, num_experts, top_k, dropout)
        else:
            self.feed_forward = FeedForward(d_model, d_ff, dropout)
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, freqs_cos, freqs_sin, mask=None):
        # Pre-Norm + 注意力 + 残差
        h = x + self.dropout(self.attention(self.attn_norm(x), freqs_cos, freqs_sin, mask))
        # Pre-Norm + 前馈 + 残差
        out = h + self.dropout(self.feed_forward(self.ff_norm(h)))
        return out


# ==================== 8. 简化版思维模块 ====================
class ThinkingModule(nn.Module):
    """轻量级推理增强模块 - 通过中间层表示增强"""
    
    def __init__(self, d_model: int, n_heads: int, n_thinking_steps: int = 3):
        super().__init__()
        self.n_thinking_steps = n_thinking_steps
        
        # 多层推理处理
        self.thinking_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=n_heads,
                dim_feedforward=d_model * 4,
                dropout=0.1,
                batch_first=True
            ) for _ in range(n_thinking_steps)
        ])
        
        self.norms = nn.ModuleList([RMSNorm(d_model) for _ in range(n_thinking_steps)])
        self.output_proj = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        current = x
        for i in range(self.n_thinking_steps):
            # 因果掩码
            seq_len = x.size(1)
            mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).bool()
            attended = self.thinking_layers[i](current, src_mask=mask)
            current = self.norms[i](current + attended)
        
        return self.output_proj(current)


# ==================== 9. 记忆模块 ====================
class MemoryModule(nn.Module):
    """记忆模块 - 简化版可学习记忆"""
    
    def __init__(self, d_model: int, memory_size: int = 10):
        super().__init__()
        self.d_model = d_model
        self.memory_size = memory_size
        
        # 可学习的记忆键值对
        self.memory_keys = nn.Parameter(torch.randn(memory_size, d_model) * 0.02)
        self.memory_values = nn.Parameter(torch.randn(memory_size, d_model) * 0.02)
    
    def forward(self, query):
        """检索记忆"""
        batch_size = query.size(0)
        keys = self.memory_keys.unsqueeze(0).expand(batch_size, -1, -1)
        values = self.memory_values.unsqueeze(0).expand(batch_size, -1, -1)
        
        # 计算注意力分数
        scores = torch.matmul(query, keys.transpose(-2, -1)) / math.sqrt(self.d_model)
        attn_weights = torch.softmax(scores, dim=-1)
        
        memory_output = torch.matmul(attn_weights, values)
        return memory_output


# ==================== 10. 主模型 ====================
class LLMModel(nn.Module):
    """改进版LLM模型"""
    
    def __init__(self, vocab_size: int, d_model: int = 512, n_heads: int = 8,
                 n_layers: int = 6, d_ff: int = 2048, max_len: int = 512, dropout: float = 0.1,
                 use_memory: bool = True, memory_size: int = 10,
                 use_thinking: bool = True, n_thinking_steps: int = 3,
                 use_moe: bool = True, num_experts: int = 4, top_k: int = 2):
        super().__init__()
        
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.d_ff = d_ff
        self.max_len = max_len
        self.use_memory = use_memory
        self.use_thinking = use_thinking
        self.use_moe = use_moe
        
        # 词嵌入
        self.embedding = nn.Embedding(vocab_size, d_model)
        
        # 预计算RoPE频率
        freqs_cos, freqs_sin = precompute_freqs_cis(d_model // n_heads, max_len * 2)
        self.register_buffer('freqs_cos', freqs_cos)
        self.register_buffer('freqs_sin', freqs_sin)
        
        # 思考模块
        if use_thinking:
            self.thinking = ThinkingModule(d_model, n_heads, n_thinking_steps)
        
        # 记忆模块
        if use_memory:
            self.memory = MemoryModule(d_model, memory_size)
        
        # Transformer层
        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads, d_ff, dropout,
                           use_moe=use_moe, num_experts=num_experts, top_k=top_k)
            for _ in range(n_layers)
        ])
        
        # 输出层
        self.fc_out = nn.Linear(d_model, vocab_size, bias=False)
        
        self.dropout = nn.Dropout(dropout)
        
        # 初始化
        self._init_weights()
    
    def _init_weights(self):
        """参数初始化"""
        for name, param in self.named_parameters():
            if param.dim() > 1:
                if 'embedding' in name or 'fc_out' in name:
                    torch.nn.init.normal_(param, mean=0.0, std=0.02)
                else:
                    torch.nn.init.normal_(param, mean=0.0, std=0.02)
    
    def forward(self, x, mask=None, return_thinking=False):
        batch_size, seq_len = x.shape
        
        # 词嵌入
        x = self.embedding(x)
        x = self.dropout(x)
        
        # 获取位置编码
        freqs_cos = self.freqs_cos[:seq_len]
        freqs_sin = self.freqs_sin[:seq_len]
        
        thinking_output = None
        
        # 思考模块
        if self.use_thinking:
            thinking_output = self.thinking(x)
            x = x + thinking_output  # 残差融合
        
        # 记忆模块
        if self.use_memory:
            memory_output = self.memory(x)
            x = x + memory_output  # 残差融合
        
        # Transformer层
        for block in self.transformer_blocks:
            x = block(x, freqs_cos, freqs_sin, mask)
        
        # 输出
        output = self.fc_out(x)
        
        if return_thinking and thinking_output is not None:
            return output, thinking_output
        return output
    
    def generate_mask(self, seq_len):
        """生成因果掩码"""
        mask = torch.tril(torch.ones(seq_len, seq_len))
        return mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq, seq]


# ==================== 11. 数据集 ====================
class ConversationDataset(Dataset):
    """对话数据集"""
    
    def __init__(self, conversations: List[Tuple[str, str]], vocab: dict, max_length: int = 128):
        self.conversations = conversations
        self.vocab = vocab
        self.max_length = max_length
        self.pad_token_id = vocab.get('<PAD>', 0)
        self.unk_token_id = vocab.get('<UNK>', 1)
    
    def __len__(self):
        return len(self.conversations)
    
    def __getitem__(self, idx):
        question, answer = self.conversations[idx]
        
        # 创建训练样本: <SOS> 问题 <SEP> <START> 答案 <END> <EOS>
        text = f"<SOS>{question}<SEP><START>{answer}<END><EOS>"
        ids = tokenizer_instance.encode(text)
        
        # 填充或截断
        ids = self.pad_or_truncate(ids)
        
        # 创建输入和目标（错位一个位置）
        input_ids = ids[:-1]
        target_ids = ids[1:]
        
        return {
            'input_ids': torch.tensor(input_ids, dtype=torch.long),
            'target_ids': torch.tensor(target_ids, dtype=torch.long)
        }
    
    def pad_or_truncate(self, ids: List[int]) -> List[int]:
        if len(ids) > self.max_length:
            return ids[:self.max_length]
        else:
            return ids + [self.pad_token_id] * (self.max_length - len(ids))


# ==================== 12. 训练器 ====================
class LLMTrainer:
    """LLM训练器"""
    
    def __init__(self, model: nn.Module, device: str = 'cpu', use_amp: bool = None):
        self.model = model.to(device)
        self.device = device
        self.criterion = nn.CrossEntropyLoss(ignore_index=0)
        self.optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.01)
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100)
        
        # 混合精度训练
        if use_amp is None:
            self.use_amp = torch.cuda.is_available()
        else:
            self.use_amp = use_amp and torch.cuda.is_available()
        
        if self.use_amp:
            self.scaler = torch.cuda.amp.GradScaler()
            print("✓ 已启用混合精度训练（AMP）")
        else:
            self.scaler = None
        
        if torch.cuda.is_available():
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.deterministic = False
        
        self.save_queue = queue.Queue()
        self.save_lock = threading.Lock()
    
    def train_epoch(self, dataloader: DataLoader, epoch: int, vocab: dict = None, config: dict = None):
        self.model.train()
        total_loss = 0
        accumulation_steps = config.get('gradient_accumulation_steps', 1) if config else 1
        log_interval = max(1, len(dataloader) // 10)  # 每10%打印一次
        
        for batch_idx, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(self.device, non_blocking=True)
            target_ids = batch['target_ids'].to(self.device, non_blocking=True)
            
            # 生成因果掩码
            mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
            
            # 前向传播
            if self.use_amp:
                with torch.cuda.amp.autocast():
                    outputs = self.model(input_ids, mask)
                    outputs = outputs.reshape(-1, outputs.size(-1))
                    target_ids = target_ids.reshape(-1)
                    loss = self.criterion(outputs, target_ids)
            else:
                outputs = self.model(input_ids, mask)
                outputs = outputs.reshape(-1, outputs.size(-1))
                target_ids = target_ids.reshape(-1)
                loss = self.criterion(outputs, target_ids)
            
            loss = loss / accumulation_steps
            
            # 反向传播
            if self.use_amp:
                self.scaler.scale(loss).backward()
                if (batch_idx + 1) % accumulation_steps == 0:
                    self.scaler.unscale_(self.optimizer)
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                    self.optimizer.zero_grad()
            else:
                loss.backward()
                if (batch_idx + 1) % accumulation_steps == 0:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                    self.optimizer.step()
                    self.optimizer.zero_grad()
            
            total_loss += loss.item() * accumulation_steps
            
            # 定期打印日志
            if (batch_idx + 1) % log_interval == 0 or batch_idx == len(dataloader) - 1:
                avg_loss = total_loss / (batch_idx + 1)
                mem_info = ""
                if torch.cuda.is_available():
                    allocated = torch.cuda.memory_allocated(self.device) / 1e9
                    reserved = torch.cuda.memory_reserved(self.device) / 1e9
                    mem_info = f", Mem: {allocated:.2f}GB/{reserved:.2f}GB"
                print(f'Epoch [{epoch}], Batch [{batch_idx + 1}/{len(dataloader)}], Loss: {avg_loss:.4f}{mem_info}')
        
        self.scheduler.step()
        return total_loss / len(dataloader)
    
    def save_model(self, filepath: str, vocab: dict, epoch: int = None):
        save_dict = {
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'vocab': vocab,
            'model_config': {
                'vocab_size': self.model.vocab_size,
                'd_model': self.model.d_model,
                'n_heads': self.model.n_heads,
                'n_layers': self.model.n_layers,
                'd_ff': self.model.d_ff,
                'max_len': self.model.max_len,
                'use_memory': self.model.use_memory,
                'use_thinking': self.model.use_thinking,
                'use_moe': self.model.use_moe
            }
        }
        if epoch is not None:
            save_dict['epoch'] = epoch
        
        torch.save(save_dict, filepath)
        print(f"模型已保存到: {filepath}")
    
    def load_model(self, filepath: str):
        checkpoint = torch.load(filepath, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        vocab = checkpoint['vocab']
        print(f"模型已从 {filepath} 加载")
        return vocab
    
    def request_save(self, filename: str, vocab: dict):
        """请求在训练过程中保存模型"""
        self.save_queue.put({'filename': filename, 'vocab': vocab})


# ==================== 13. 对话生成 ====================
class LLMChat:
    """LLM对话"""
    
    def __init__(self, model: nn.Module, vocab: dict, device: str = 'cpu', max_length: int = 128,
                 show_thinking: bool = True):
        self.model = model.to(device)
        self.model.eval()
        self.vocab = vocab
        self.device = device
        self.max_length = max_length
        self.id_to_token = {v: k for k, v in vocab.items()}
        self.show_thinking = show_thinking
        self.conversation_history = []
        self.max_history = 10
        
        # 同步全局tokenizer的vocab（关键！）
        tokenizer_instance.char_to_id = vocab
        tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()}
    
    def generate_response(self, question: str, max_length: int = 100, temperature: float = 0.8,
                         include_thinking: bool = False, use_beam_search: bool = False):
        """生成回答"""
        self.model.eval()
        
        # 使用与训练数据一致的格式: <SOS>问题<SEP><START>
        input_text = f"<SOS>{question}<SEP><START>"
        
        input_ids = tokenizer_instance.encode(input_text)
        
        if use_beam_search:
            return self._generate_with_beam_search(input_ids, max_length, temperature, include_thinking)
        else:
            return self._generate_with_sampling(input_ids, max_length, temperature, include_thinking)
    
    def _generate_with_sampling(self, input_ids: List[int], max_length: int, temperature: float, include_thinking: bool):
        """使用采样生成文本，带重复惩罚"""
        input_tensor = torch.tensor([input_ids], dtype=torch.long).to(self.device)
        generated = input_ids.copy()
        
        with torch.no_grad():
            for step in range(max_length):
                seq = input_tensor[:, -self.model.max_len:]
                
                # 前向传播
                outputs = self.model(seq)
                
                # 获取最后一个时间步的预测
                next_token_logits = outputs[0, -1, :] / temperature
                
                # === 重复惩罚 ===
                # 1. 对最近生成的token进行惩罚（更强）
                recent_tokens = generated[-30:]  # 最近30个token
                for token_id in set(recent_tokens):
                    count = recent_tokens.count(token_id)
                    if count > 1:
                        next_token_logits[token_id] -= 3.0 * count
                
                # 2. 阻止2-gram重复（更强）
                if len(generated) >= 2:
                    last_bigram = (generated[-2], generated[-1])
                    for i in range(len(generated) - 2):
                        if (generated[i], generated[i+1]) == last_bigram:
                            if i + 2 < len(generated):
                                next_token_logits[generated[i+2]] -= 10.0
                
                # 3. 阻止3-gram重复
                if len(generated) >= 3:
                    last_trigram = (generated[-3], generated[-2], generated[-1])
                    for i in range(len(generated) - 3):
                        if (generated[i], generated[i+1], generated[i+2]) == last_trigram:
                            if i + 3 < len(generated):
                                next_token_logits[generated[i+3]] -= 15.0
                
                # 4. 阻止完全重复句子（检测连续重复模式）
                decoded_so_far = tokenizer_instance.decode(generated[len(input_ids):])
                if self._has_repeated_pattern(decoded_so_far):
                    break
                
                # Top-k采样
                next_token = self._top_k_sampling(next_token_logits, top_k=30)
                
                # 检查结束标记
                if next_token == self.vocab.get('<END>', -1):
                    break
                
                generated.append(next_token)
                
                # 更新输入
                new_token = torch.tensor([[next_token]], device=self.device)
                input_tensor = torch.cat([input_tensor, new_token], dim=1)
                
                if len(generated) > max_length:
                    break
        
        # 解码
        generated_text = tokenizer_instance.decode(generated[len(input_ids):])
        generated_text = self._post_process_response(generated_text)
        
        return generated_text
    
    def _has_repeated_pattern(self, text: str) -> bool:
        """检测文本中是否有重复3次以上的模式"""
        if len(text) < 6:
            return False
        # 检测连续重复2次以上的子串（长度>=2）
        for length in range(2, min(12, len(text)//2 + 1)):
            for i in range(len(text) - 2*length + 1):
                pattern = text[i:i+length]
                if text[i:i+2*length] == pattern * 2:
                    return True
        return False
    
    def _generate_with_beam_search(self, input_ids: List[int], max_length: int, temperature: float,
                                   include_thinking: bool, beam_width: int = 3):
        """使用束搜索生成文本"""
        beams = [(input_ids.copy(), 0.0)]
        
        for _ in range(max_length):
            candidates = []
            
            for seq, score in beams:
                input_tensor = torch.tensor([seq], dtype=torch.long).to(self.device)
                
                with torch.no_grad():
                    outputs = self.model(input_tensor[:, -self.model.max_len:])
                    next_token_logits = outputs[0, -1, :] / temperature
                    next_token_probs = torch.softmax(next_token_logits, dim=-1)
                    
                    top_k_probs, top_k_indices = torch.topk(next_token_probs, beam_width)
                    
                    for i in range(beam_width):
                        token_id = top_k_indices[i].item()
                        token_prob = top_k_probs[i].item()
                        new_seq = seq + [token_id]
                        new_score = score + math.log(token_prob)
                        candidates.append((new_seq, new_score))
            
            candidates.sort(key=lambda x: x[1], reverse=True)
            beams = candidates[:beam_width]
            
            # 检查是否结束
            if all(beam[0][-1] == self.vocab.get('<END>', -1) for beam in beams):
                break
        
        best_seq = beams[0][0]
        generated_text = tokenizer_instance.decode(best_seq[len(input_ids):])
        generated_text = self._post_process_response(generated_text)
        
        return generated_text
    
    def _top_k_sampling(self, logits, top_k=50):
        """Top-k采样"""
        if top_k > 0:
            top_k = min(top_k, logits.size(-1))
            indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
            logits = logits.masked_fill(indices_to_remove, float('-inf'))
        
        probs = torch.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        return next_token.item()
    
    def _post_process_response(self, response: str) -> str:
        """后处理响应"""
        response = response.replace('<SOS>', '').replace('<EOS>', '')
        response = response.replace('<START>', '').replace('<END>', '')
        response = response.replace('<SEP>', '')
        response = response.strip()
        
        # 去除重复片段（如 "我是小狸，我是小狸"）
        response = self._remove_repeated_phrases(response)
        
        # 确保句子完整性
        if response and response[-1] not in ['。', '！', '？', '.', '!', '?', '\n']:
            last_punct = max(response.rfind(p) for p in ['。', '！', '？', '.', '!', '?', '\n'])
            if last_punct != -1:
                response = response[:last_punct+1]
        
        return response if response else '...'
    
    def _remove_repeated_phrases(self, text: str) -> str:
        """去除文本中的重复短语"""
        if len(text) < 4:
            return text
        
        # 检测并去除重复2次以上的子串
        for length in range(4, min(20, len(text)//2 + 1)):
            i = 0
            while i <= len(text) - 2*length:
                pattern = text[i:i+length]
                if text[i:i+2*length] == pattern * 2:
                    # 找到重复，保留一个
                    text = text[:i] + pattern + text[i+2*length:]
                    # 继续检查同一位置
                else:
                    i += 1
        
        # 清理多余的标点
        text = re.sub(r'([。！？.!?])\1+', r'\1', text)
        
        return text.strip()
    
    def clear_history(self):
        """清除对话历史"""
        self.conversation_history = []
    
    def get_history(self):
        """获取对话历史"""
        return self.conversation_history
    
    def toggle_thinking(self):
        """切换是否显示思考过程"""
        self.show_thinking = not self.show_thinking
        return self.show_thinking


# ==================== 14. 数据加载 ====================
def load_training_data(data_path: str) -> List[Tuple[str, str]]:
    """加载训练数据"""
    possible_paths = [
        data_path,
        os.path.join(os.path.dirname(__file__), data_path),
        os.path.join(os.getcwd(), data_path)
    ]
    
    for path in possible_paths:
        try:
            if os.path.exists(path):
                with open(path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    
                    if isinstance(data, list):
                        if data and isinstance(data[0], dict):
                            if 'question' in data[0] and 'answer' in data[0]:
                                return [(item['question'], item['answer']) for item in data]
                            elif 'input' in data[0] and 'output' in data[0]:
                                return [(item['input'], item['output']) for item in data]
                            elif 'prompt' in data[0] and 'response' in data[0]:
                                return [(item['prompt'], item['response']) for item in data]
                            elif 'instruction' in data[0] and 'output' in data[0]:
                                return [(item['instruction'], item['output']) for item in data]
                            elif 'conversations' in data[0]:
                                conversations = []
                                for item in data:
                                    convs = item['conversations']
                                    for i in range(0, len(convs)-1, 2):
                                        if i+1 < len(convs):
                                            conversations.append((convs[i]['value'], convs[i+1]['value']))
                                return conversations
                            elif 'human' in data[0] and 'assistant' in data[0]:
                                return [(item['human'], item['assistant']) for item in data]
                            else:
                                print(f"未知的数据格式，字段包括: {list(data[0].keys())}")
                                return []
                        elif data and isinstance(data[0], list):
                            return [(str(item[0]), str(item[1])) for item in data if len(item) >= 2]
                        else:
                            return []
                    elif isinstance(data, dict):
                        if 'data' in data and isinstance(data['data'], list):
                            if data['data'] and 'question' in data['data'][0] and 'answer' in data['data'][0]:
                                return [(item['question'], item['answer']) for item in data['data']]
                        conversations = []
                        if 'content' in data:
                            content = data['content']
                            if isinstance(content, list):
                                for item in content:
                                    if isinstance(item, dict):
                                        if 'q' in item and 'a' in item:
                                            conversations.append((str(item['q']), str(item['a'])))
                                        elif 'question' in item and 'answer' in item:
                                            conversations.append((str(item['question']), str(item['answer'])))
                        return conversations
                    return []
        except json.JSONDecodeError:
            print(f"错误：{path} 不是有效的JSON文件")
            return []
        except Exception as e:
            print(f"加载数据失败（{path}）：{str(e)}")
            continue
    
    print(f"警告：未找到数据文件 {data_path}")
    return []


# ==================== 15. 主函数 ====================
def main():
    print("=" * 50)
    print("LLM对话AI训练程序 - 改进版")
    print("=" * 50)
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    # 选择训练数据文件
    print("\n可用的训练数据文件（在当前目录）：")
    import glob
    json_files = glob.glob("*.json")
    if json_files:
        for i, f in enumerate(json_files, 1):
            print(f"  {i}. {f}")
    else:
        print("  未找到JSON数据文件")
    
    data_path = input("\n请输入训练数据文件路径（或直接回车使用默认: training_data.json）: ").strip()
    data_path = data_path.strip('"').strip("'")
    if not data_path:
        data_path = 'training_data.json'
    
    config = {
        'data_path': data_path,
        'model_path': 'llm_model.pth',
        'batch_size': 32 if torch.cuda.is_available() else 4,
        'gradient_accumulation_steps': 2,
        'epochs': 100,
        'd_model': 512,
        'n_heads': 8,
        'n_layers': 8,
        'd_ff': 2048,
        'max_length': 128,
        'dropout': 0.1,
        'use_amp': torch.cuda.is_available(),
        'generation_config': {
            'top_k': 50,
            'top_p': 0.9,
            'temperature': 0.8
        }
    }
    
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"🔥 检测到CUDA设备: {gpu_name}")
        print(f"  💾 显存总量: {gpu_mem:.2f}GB")
        print(f"  📦 批次大小: {config['batch_size']}")
        print(f"  ⚡ 混合精度训练: {'启用' if config['use_amp'] else '禁用'}")
    else:
        print("⚠ 未检测到CUDA设备，将使用CPU训练")
    
    # 加载数据
    print(f"\n加载训练数据: {config['data_path']}")
    conversations = load_training_data(config['data_path'])
    if not conversations:
        print(f"\n错误：无法加载训练数据 {config['data_path']}")
        return
    print(f"加载了 {len(conversations)} 条对话数据")
    
    # 构建词汇表
    print("\n构建词汇表...")
    vocab = build_vocabulary(conversations)
    print(f"词汇表大小: {len(vocab)}")
    
    # 创建数据集
    dataset = ConversationDataset(conversations, vocab, config['max_length'])
    dataloader = DataLoader(
        dataset,
        batch_size=config['batch_size'],
        shuffle=True,
        num_workers=4 if torch.cuda.is_available() else 0,
        pin_memory=True if torch.cuda.is_available() else False,
        persistent_workers=True if torch.cuda.is_available() else False
    )
    print(f"数据集大小: {len(dataset)}")
    print(f"DataLoader批次数量: {len(dataloader)}")
    
    # 询问是否继续训练
    print("\n" + "=" * 50)
    print("训练模式选择")
    print("=" * 50)
    
    pth_files = glob.glob("*.pth")
    if pth_files:
        for i, f in enumerate(pth_files, 1):
            print(f"  {i}. {f}")
    else:
        print("  未找到模型文件")
    
    resume_choice = input("\n是否从已有模型继续训练？(y/n，默认n): ").strip().lower()
    resume_training = resume_choice == 'y' or resume_choice == 'yes'
    
    start_epoch = 1
    resume_model_path = None
    if resume_training and pth_files:
        print("\n可用的模型文件：")
        for i, f in enumerate(pth_files, 1):
            print(f"  {i}. {f}")
        
        model_choice = input(f"\n请输入要加载的模型编号（1-{len(pth_files)}）: ").strip()
        try:
            model_idx = int(model_choice) - 1
            if 0 <= model_idx < len(pth_files):
                resume_model_path = pth_files[model_idx]
                print(f"\n选择加载模型: {resume_model_path}")
        except (ValueError, IndexError):
            print("\n无效的选择，将重新开始训练")
            resume_training = False
    
    # 训练前参数配置
    print("\n" + "=" * 50)
    print("训练参数配置")
    print("=" * 50)
    
    memory_choice = input("是否启用记忆功能？(Y/n，默认Y): ").strip().lower()
    use_memory = memory_choice != 'n' and memory_choice != 'no'
    
    thinking_choice = input("是否启用深度思考？(Y/n，默认Y): ").strip().lower()
    use_thinking = thinking_choice != 'n' and thinking_choice != 'no'
    
    moe_choice = input("是否启用MoE？(Y/n，默认Y): ").strip().lower()
    use_moe = moe_choice != 'n' and moe_choice != 'no'
    
    if use_thinking:
        try:
            n_thinking_steps = int(input("思考步骤数 (默认3): ").strip() or "3")
        except ValueError:
            n_thinking_steps = 3
    else:
        n_thinking_steps = 1
    
    print(f"\n当前配置:")
    print(f"记忆功能: {'已启用' if use_memory else '未启用'}")
    print(f"深度思考: {'已启用' if use_thinking else '未启用'}")
    print(f"MoE: {'已启用' if use_moe else '未启用'}")
    print(f"思考步骤数: {n_thinking_steps}")
    
    confirm = input("\n是否使用以上配置开始训练？(Y/n，默认Y): ").strip().lower()
    if confirm == 'n' or confirm == 'no':
        print("训练已取消")
        return
    
    # 创建模型
    print("\n创建LLM模型...")
    model = LLMModel(
        vocab_size=len(vocab),
        d_model=config['d_model'],
        n_heads=config['n_heads'],
        n_layers=config['n_layers'],
        d_ff=config['d_ff'],
        max_len=config['max_length'],
        dropout=config['dropout'],
        use_memory=use_memory,
        memory_size=10,
        use_thinking=use_thinking,
        n_thinking_steps=n_thinking_steps,
        use_moe=use_moe,
        num_experts=4 if use_moe else 1,
        top_k=2 if use_moe else 1
    )
    print(f"模型配置: d_model={config['d_model']}, n_layers={config['n_layers']}, n_heads={config['n_heads']}")
    print(f"记忆功能: {'已启用' if model.use_memory else '未启用'}")
    print(f"深度思考: {'已启用' if model.use_thinking else '未启用'}")
    print(f"MoE: {'已启用' if model.use_moe else '未启用'}")
    
    # 创建训练器
    trainer = LLMTrainer(model, device, use_amp=config['use_amp'])
    print(f"\n🎮 训练器配置:")
    print(f"   - 混合精度: {'✓ 启用 (FP16)' if config['use_amp'] else '✗ 禁用 (FP32)'}")
    print(f"   - 梯度累积: {config['gradient_accumulation_steps']} 步")
    
    # 加载已有模型
    if resume_training and resume_model_path:
        try:
            print(f"\n加载模型参数: {resume_model_path}")
            checkpoint = torch.load(resume_model_path, map_location=device)
            
            if 'epoch' in checkpoint:
                start_epoch = checkpoint['epoch'] + 1
                print(f"从第 {checkpoint['epoch']} 轮继续")
            else:
                import re
                epoch_match = re.search(r'epoch[_]?(\d+)', resume_model_path)
                if epoch_match:
                    start_epoch = int(epoch_match.group(1)) + 1
                    print(f"从文件名推断：从第 {start_epoch} 轮继续训练")
                else:
                    start_epoch = 1
            
            trainer.load_model(resume_model_path)
            print(f"✓ 模型参数加载成功")
            
            if start_epoch > 1:
                config['epochs'] = config['epochs'] + start_epoch - 1
                print(f"训练轮数调整为: {config['epochs']} 轮")
        except Exception as e:
            print(f"✗ 加载模型失败: {e}")
            resume_training = False
            start_epoch = 1
    
    # 训练
    print("\n" + "="*60)
    print("🚀 开始训练！")
    print("="*60)
    start_time = time.time()
    
    try:
        for epoch in range(start_epoch, config['epochs'] + 1):
            epoch_start = time.time()
            print(f"\n{'='*60}")
            print(f"📊 Epoch {epoch}/{config['epochs']} 开始")
            print(f"{'='*60}")
            
            avg_loss = trainer.train_epoch(dataloader, epoch, vocab, config)
            epoch_time = time.time() - epoch_start
            
            print(f"\n{'='*60}")
            print(f"📈 Epoch {epoch} 完成 | 耗时: {epoch_time:.1f}秒")
            print(f'   📉 Loss: {avg_loss:.4f}')
            if torch.cuda.is_available():
                print(f"   💾 显存峰值: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")
                torch.cuda.reset_peak_memory_stats()
            print(f"{'='*60}\n")
            
            # 每10个epoch保存一次
            if epoch % 10 == 0:
                save_path = f'llm_model_epoch_{epoch}.pth'
                print(f"💾 保存检查点: {save_path}")
                trainer.save_model(save_path, vocab, epoch)
        
        # 训练完成
        training_time = time.time() - start_time
        print("\n" + "="*60)
        print("🎉 训练完成！")
        print("="*60)
        print(f"⏱️  总耗时: {training_time:.1f}秒 ({training_time/60:.1f}分钟)")
        print("="*60)
        
        # 保存最终模型
        print(f"💾 保存最终模型: {config['model_path']}")
        trainer.save_model(config['model_path'], vocab, config['epochs'])
    
    except KeyboardInterrupt:
        print("\n检测到中断信号，正在保存模型...")
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        trainer.save_model(f'llm_model_interrupted_{timestamp}.pth', vocab)
        print("模型已保存，训练终止")
    
    # 测试
    print("\n" + "=" * 50)
    print("测试模型")
    print("=" * 50)
    
    llm_chat = LLMChat(model, vocab, device, config['max_length'])
    
    test_questions = ["你好", "你是谁", "谢谢", "再见"]
    for q in test_questions:
        response = llm_chat.generate_response(q)
        print(f"\n问题: {q}")
        print(f"回答: {response}")
    
    # 交互模式
    print("\n" + "=" * 50)
    print("进入交互对话模式")
    print("=" * 50)
    print("可用命令：")
    print("  quit or 退出  - 退出对话")
    print("  clear or 清除  - 清除对话历史")
    print("=" * 50)
    
    while True:
        user_input = input("\n你: ").strip()
        if user_input.lower() in ['quit', '退出', '再见']:
            print("再见！")
            break
        if user_input.lower() in ['clear', '清除']:
            llm_chat.clear_history()
            print("对话历史已清除")
            continue
        if not user_input:
            continue
        
        response = llm_chat.generate_response(user_input)
        print(f"AI: {response}")


if __name__ == "__main__":
    main()