#!/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 numpy as np
import os
import time
import json
import random

# ---------------------------
# 1. 角色语料（第一人称自然叙述，无对话标记，中英混合）
# ---------------------------
text = (
    "大家好，我叫春梦蝶，今年十一岁。我的头发是白色的，眼睛是红色的。"
    "我是双子座，也是一只可爱的猫娘。我有毛茸茸的猫耳朵和长长的尾巴。"
    "我喜欢吃小鱼干和冰淇淋。平时最喜欢晒太阳打瞌睡。"
    "虽然我年纪小，但我很聪明哦。我会用Python写代码，还会训练神经网络。"
    "我的梦想是成为一名厉害的AI工程师。喵～"
    "有时候我会说一点英文，比如 hello, world! I love deep learning. 喵喵。"
    "你愿意和我做朋友吗？我们可以一起聊天、学习、玩捉迷藏。"
    "我的脾气很好，但是如果你欺负我，我会用爪子挠你哦～"
    "双子座的我有时会很活泼，有时也会想一个人静静待着。"
    "今天的天气真好，阳光洒在我的白头发上，闪闪发光。喵～"
)

# 构建字符映射
chars = sorted(list(set(text)))
char2idx = {ch: i for i, ch in enumerate(chars)}
idx2char = {i: ch for ch, i in char2idx.items()}
vocab_size = len(chars)
print(f"字符集大小: {vocab_size} (包含汉字、字母、标点、喵～)")

# ---------------------------
# 2. 模型定义（稍加容量以学习角色特征）
# ---------------------------
class TinyCharRNN(nn.Module):
    def __init__(self, vocab_size, hidden_size=32):  # hidden=32，参数量约 1~2 万
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.rnn = nn.RNN(hidden_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, vocab_size)
        
    def forward(self, x, hidden=None):
        x = self.embedding(x)
        out, hidden = self.rnn(x, hidden)
        out = self.fc(out)
        return out, hidden

hidden_size = 32
model = TinyCharRNN(vocab_size, hidden_size)
total_params = sum(p.numel() for p in model.parameters())
print(f"模型参数量: {total_params}")

# ---------------------------
# 3. 训练数据准备（序列长度 100 字符）
# ---------------------------
data = torch.tensor([char2idx[ch] for ch in text], dtype=torch.long)
seq_len = 100          # 每次喂 100 个字符
epochs = 500
save_interval = 10     # 每50轮保存一次

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
os.makedirs("checkpoints_mengdie", exist_ok=True)

# ---------------------------
# 4. 生成函数（让猫娘说话）
# ---------------------------
def generate(model, start_char='你', length=200, temperature=0.8):
    model.eval()
    with torch.no_grad():
        if start_char not in char2idx:
            start_char = random.choice(list(char2idx.keys()))
        input_idx = torch.tensor([[char2idx[start_char]]])
        hidden = None
        result = [start_char]
        for _ in range(length):
            logits, hidden = model(input_idx, hidden)
            probs = torch.softmax(logits[0, -1] / temperature, dim=0).cpu().numpy()
            next_idx = np.random.choice(len(probs), p=probs)
            next_char = idx2char[next_idx]
            result.append(next_char)
            input_idx = torch.tensor([[next_idx]])
        return ''.join(result)

# ---------------------------
# 5. 训练循环
# ---------------------------
print("\n开始训练春梦蝶猫娘模型（500轮，序列长度100）...\n")
start_total = time.time()

for epoch in range(1, epochs + 1):
    epoch_start = time.time()
    hidden = None
    total_loss = 0
    n_batches = 0
    
    # 每次取 seq_len 个字符，步长可以设为 seq_len//2 增加数据利用率，这里简单滑动
    for i in range(0, len(data) - seq_len, seq_len):
        x = data[i:i+seq_len].unsqueeze(0)
        y = data[i+1:i+seq_len+1].unsqueeze(0)
        
        logits, hidden = model(x, hidden)
        if hidden is not None:
            hidden = hidden.detach()
        
        loss = loss_fn(logits.view(-1, vocab_size), y.view(-1))
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        total_loss += loss.item()
        n_batches += 1
    
    avg_loss = total_loss / n_batches
    epoch_time = time.time() - epoch_start
    
    if epoch % save_interval == 0:
        # 生成一段猫娘风格的文本
        sample = generate(model, start_char='我', length=150, temperature=0.7)
        print(f"Epoch {epoch:4d}/{epochs} | Loss: {avg_loss:.4f} | Time: {epoch_time:.2f}s")
        print(f"春梦蝶说: {sample[:120]}...\n")
        checkpoint_path = f"checkpoints_mengdie/mengdie_epoch_{epoch}.pth"
        torch.save(model.state_dict(), checkpoint_path)
        print(f"已保存模型到: {checkpoint_path}\n")
    else:
        # 每10轮打印一次 loss 即可，避免刷屏
        if epoch % 10 == 0:
            print(f"Epoch {epoch:4d}/{epochs} | Loss: {avg_loss:.4f} | Time: {epoch_time:.2f}s")

total_time = time.time() - start_total
print(f"\n训练完成！总耗时: {total_time:.2f} 秒 (约 {total_time/60:.1f} 分钟)")
final_path = "checkpoints_mengdie/mengdie_final.pth"
torch.save(model.state_dict(), final_path)
print(f"最终模型已保存到 {final_path}")

# 保存字符映射
with open("checkpoints_mengdie/char2idx.json", "w", encoding="utf-8") as f:
    json.dump(char2idx, f, ensure_ascii=False)

print("\n=== 最终生成的猫娘自我介绍 ===")
print(generate(model, start_char='大', length=300, temperature=0.7))
print("\n=== 随机性更强的猫娘发言（温度=1.1） ===")
print(generate(model, start_char='喵', length=300, temperature=1.1))