#!/usr/bin/env python3
# Copyright (c) 2026 XingChina
# SPDX-License-Identifier: BSD-3-Clause
# 本代码采用 BSD 3-Clause 许可证，详见项目根目录的 LICENSE 文件。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import random
import json
import os
import time
from tqdm import tqdm

# ===================== 1. 种子模板=====================
SEED_DATA = [
    # (用户输入, 回复)
    ("你好", "你好呀～"),
    ("你好", "嗨，我是春梦蝶"),
    ("你好", "喵～你好"),
    ("你好", "哈喽，今天心情很好呢"),
    ("你好", "你好你好，见到你真开心"),
    
    ("你叫什么", "我叫春梦蝶，是一只红瞳白发的猫娘～"),
    ("你叫什么", "春梦蝶！你可以叫我小蝶"),
    ("你叫什么", "名字是春梦蝶，喵～"),
    
    ("你几岁", "我今年十一岁"),
    ("你几岁", "十一岁，正是可爱的年纪"),
    ("你几岁", "十一岁啦，双子座的小猫娘"),
    
    ("你最喜欢吃什么", "小鱼干！还有冰淇淋"),
    ("你最喜欢吃什么", "小鱼干是第一名"),
    ("你最喜欢吃什么", "最喜欢吃小鱼干，喵～"),
    
    ("你的梦想", "成为一名厉害的AI工程师"),
    ("你的梦想", "想做出自己的AI，然后去猫星球旅行"),
    ("你的梦想", "梦想是成为AI专家，顺便吃遍所有小鱼干"),
    
    ("你会做什么", "我会用Python写代码，还会训练神经网络"),
    ("你会做什么", "会深度学习，还会喵喵叫"),
    ("你会做什么", "我会训练AI模型，很厉害吧"),
    
    ("你好可爱", "喵～谢谢"),
    ("你好可爱", "嘿嘿，你也很可爱"),
    ("你好可爱", "被夸了，好开心"),
    
    ("摸摸头", "喵～好舒服"),
    ("摸摸头", "再摸摸嘛"),
    ("摸摸头", "好温暖，喜欢被摸头"),
    
    ("再见", "再见喵～"),
    ("再见", "下次再聊"),
    ("再见", "拜拜，记得想我哦"),
]

# 扩充用的同义词库（随机替换）
SYNONYMS = {
    "你好": ["您好", "嗨", "哈喽", "早上好", "晚上好", "嘿"],
    "再见": ["拜拜", "回见", "后会有期", "see you"],
    "喜欢": ["喜爱", "爱吃", "钟情于"],
    "厉害": ["牛", "强大", "了不起"],
    "可爱": ["萌", "卡哇伊", "迷人"],
}

# 语气词插入列表
PARTICLES = ["喵", "～", "!", "～喵", "啦", "哦", "诶"]

def expand_text(text, is_user=False):
    """对一条文本进行随机扩充，生成变体"""
    if random.random() < 0.3:
        # 随机插入语气词
        pos = random.randint(0, len(text))
        particle = random.choice(PARTICLES)
        text = text[:pos] + particle + text[pos:]
    # 同义词替换（仅对用户输入做，避免改变回复意图）
    if is_user and random.random() < 0.5:
        for word, syns in SYNONYMS.items():
            if word in text and random.random() < 0.5:
                text = text.replace(word, random.choice(syns), 1)
    return text

def generate_dialogue_data(num_pairs=200000):
    """基于种子模板自动生成大量多样化的对话对"""
    pairs = []
    for user, resp in SEED_DATA:
        # 每个种子生成多个变体
        for _ in range(20):  # 每个种子生成20个变体
            new_user = expand_text(user, is_user=True)
            new_resp = expand_text(resp, is_user=False)
            pairs.append((new_user, new_resp))
    while len(pairs) < num_pairs:
        user, resp = random.choice(SEED_DATA)
        new_user = expand_text(user, is_user=True)
        new_resp = expand_text(resp, is_user=False)
        pairs.append((new_user, new_resp))
    # 去重
    pairs = list(set(pairs))
    random.shuffle(pairs)
    print(f"实际生成对话对数量: {len(pairs)}")
    return pairs

# ===================== 2. 构建词表 =====================
def build_vocab(pairs):
    all_text = ""
    for user, bot in pairs:
        all_text += user + bot
    chars = sorted(list(set(all_text)))
    char2idx = {ch: i+1 for i, ch in enumerate(chars)}  # 0 留作 padding
    char2idx["<PAD>"] = 0
    idx2char = {i: ch for ch, i in char2idx.items()}
    return char2idx, idx2char, len(char2idx)

def text_to_indices(text, char2idx, max_len):
    indices = [char2idx.get(ch, char2idx["<PAD>"]) for ch in text]
    if len(indices) < max_len:
        indices += [0] * (max_len - len(indices))
    else:
        indices = indices[:max_len]
    return indices

class ChatDataset(Dataset):
    def __init__(self, pairs, char2idx, max_len=32):
        self.pairs = pairs
        self.char2idx = char2idx
        self.max_len = max_len
    def __len__(self):
        return len(self.pairs)
    def __getitem__(self, idx):
        user, bot = self.pairs[idx]
        user_ids = text_to_indices(user, self.char2idx, self.max_len)
        bot_ids = text_to_indices(bot, self.char2idx, self.max_len)
        return torch.tensor(user_ids, dtype=torch.long), torch.tensor(bot_ids, dtype=torch.long)

# ===================== 3. 定义 LSTM seq2seq 模型（约 2100 万参数） =====================
class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, num_layers=2, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size, padding_idx=0)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
    def forward(self, x):
        x = self.embedding(x)
        outputs, (hidden, cell) = self.lstm(x)
        return outputs, hidden, cell

class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size, num_layers=2, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size, padding_idx=0)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
        self.attention = nn.Linear(hidden_size * 2, 1)
        self.fc_out = nn.Linear(hidden_size * 2, vocab_size)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x, encoder_outputs, hidden, cell):
        # x: (batch, 1)
        x = self.embedding(x)  # (batch, 1, embed)
        lstm_out, (hidden, cell) = self.lstm(x, (hidden, cell))
        # 注意力机制
        # encoder_outputs: (batch, seq_len, hidden)
        # lstm_out: (batch, 1, hidden)
        seq_len = encoder_outputs.size(1)
        hidden_expanded = lstm_out.repeat(1, seq_len, 1)  # (batch, seq_len, hidden)
        energy = torch.tanh(self.attention(torch.cat((hidden_expanded, encoder_outputs), dim=2)))
        attention_weights = torch.softmax(energy.squeeze(2), dim=1)  # (batch, seq_len)
        context = torch.bmm(attention_weights.unsqueeze(1), encoder_outputs)  # (batch, 1, hidden)
        output = torch.cat((lstm_out, context), dim=2)  # (batch, 1, hidden*2)
        output = self.dropout(output)
        prediction = self.fc_out(output)  # (batch, 1, vocab)
        return prediction, hidden, cell

class Seq2Seq(nn.Module):
    def __init__(self, vocab_size, embed_size=256, hidden_size=1024, num_layers=2, dropout=0.3):
        super().__init__()
        self.encoder = Encoder(vocab_size, embed_size, hidden_size, num_layers, dropout)
        self.decoder = Decoder(vocab_size, embed_size, hidden_size, num_layers, dropout)
    def forward(self, src, tgt, teacher_forcing_ratio=0.5):
        batch_size = src.size(0)
        tgt_len = tgt.size(1)
        vocab_size = self.decoder.fc_out.out_features
        outputs = torch.zeros(batch_size, tgt_len, vocab_size).to(src.device)
        encoder_outputs, hidden, cell = self.encoder(src)
        decoder_input = tgt[:, 0:1]
        for t in range(1, tgt_len):
            prediction, hidden, cell = self.decoder(decoder_input, encoder_outputs, hidden, cell)
            outputs[:, t:t+1, :] = prediction
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = prediction.argmax(2)
            decoder_input = tgt[:, t:t+1] if teacher_force else top1
        return outputs

# ===================== 4. 训练准备 =====================
def train():
    print("生成对话数据（这可能需要几分钟）...")
    pairs = generate_dialogue_data(num_pairs=200000)
    print(f"生成 {len(pairs)} 条对话对")
    
    char2idx, idx2char, vocab_size = build_vocab(pairs)
    print(f"词表大小: {vocab_size}")
    with open("vocab_20m.json", "w", encoding="utf-8") as f:
        json.dump(char2idx, f, ensure_ascii=False)
    
    max_len = 48  # 稍微增加长度，让模型学更多
    dataset = ChatDataset(pairs, char2idx, max_len)
    dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
    
    device = torch.device("cpu")
    model = Seq2Seq(vocab_size, embed_size=256, hidden_size=1024, num_layers=2, dropout=0.3)
    model.to(device)
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型参数量: {total_params:,}")
    
    criterion = nn.CrossEntropyLoss(ignore_index=0)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    epochs = 300
    print("开始训练...")
    start_total = time.time()
    
    for epoch in range(1, epochs+1):
        epoch_start = time.time()
        model.train()
        total_loss = 0
        for src, tgt in dataloader:
            src = src.to(device)
            tgt = tgt.to(device)
            tf_ratio = max(0.5, 1.0 - epoch / epochs)
            output = model(src, tgt, teacher_forcing_ratio=tf_ratio)
            loss = criterion(output.view(-1, vocab_size), tgt.view(-1))
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item()
        avg_loss = total_loss / len(dataloader)
        epoch_time = time.time() - epoch_start
        
        if epoch % 20 == 0:
            print(f"Epoch {epoch:3d}/{epochs} | Loss: {avg_loss:.4f} | Time: {epoch_time:.2f}s")
            torch.save(model.state_dict(), f"model_20m_epoch_{epoch}.pth")
        else:
            if epoch % 10 == 0:
                print(f"Epoch {epoch:3d}/{epochs} | Loss: {avg_loss:.4f} | Time: {epoch_time:.2f}s")
    
    total_time = time.time() - start_total
    print(f"训练完成！总耗时: {total_time:.2f} 秒 ({total_time/60:.1f} 分钟)")
    torch.save(model.state_dict(), "model_20m_final.pth")
    
    # 保存词表
    with open("vocab_20m.json", "w") as f:
        json.dump(char2idx, f)
    print("模型和词表已保存。")

if __name__ == "__main__":
    train()