import math
import random
import time
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

TEXT = (
    "起初只有一个简单的故事。"
    "它不长，但足够教会模型去猜下一个字符。"
    "我们会在这段文字上训练一个很小的语言模型，并观察它学习。"
    "实践胜过空谈，所以先从小处开始，不断迭代。"
    "很快我们会用更大的数据替换它。"
    "现在先保持小而清晰。"
)

@dataclass
class Config:
    # 序列长度（模型一次看到的字符数）
    seq_len: int = 64
    # 批大小
    batch_size: int = 32
    # 字符嵌入维度
    embed_dim: int = 64
    # GRU 隐藏层维度
    hidden_dim: int = 128
    # GRU 层数
    num_layers: int = 1
    # 学习率
    lr: float = 2e-3
    # 训练步数
    steps: int = 800
    # 采样频率
    sample_every: int = 200
    # 采样长度
    sample_len: int = 200
    # 采样温度（越高越随机）
    temperature: float = 0.9


class CharRNN(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers):
        super().__init__()
        # 字符索引 -> 向量
        self.embed = nn.Embedding(vocab_size, embed_dim)
        # GRU 用于处理序列上下文
        self.rnn = nn.GRU(
            embed_dim,
            hidden_dim,
            num_layers=num_layers,
            batch_first=True,
        )
        # 隐藏状态 -> 下一个字符的 logits
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, h=None):
        # x: (batch, seq_len)
        x = self.embed(x)
        # out: (batch, seq_len, hidden_dim)
        out, h = self.rnn(x, h)
        # logits: (batch, seq_len, vocab_size)
        logits = self.fc(out)
        return logits, h


def build_vocab(text):
    # 构建字符表与双向映射
    chars = sorted(set(text))
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for ch, i in stoi.items()}
    return chars, stoi, itos


def encode(text, stoi):
    # 文本 -> 索引序列
    return torch.tensor([stoi[ch] for ch in text], dtype=torch.long)


def get_batch(data, cfg):
    # 随机抽取若干段序列，生成输入 x 和目标 y（y 是 x 向后移一位）
    max_start = len(data) - cfg.seq_len - 1
    starts = torch.randint(0, max_start, (cfg.batch_size,))
    x = torch.stack([data[s : s + cfg.seq_len] for s in starts])
    y = torch.stack([data[s + 1 : s + cfg.seq_len + 1] for s in starts])
    return x, y


def sample(model, stoi, itos, device, cfg, start="起初"):
    # 从给定起始串开始，按概率逐字符采样
    model.eval()
    # 如果起始字符不在词表中，使用文本开头的字符兜底
    safe_start = "".join(ch for ch in start if ch in stoi)
    if not safe_start:
        safe_start = TEXT[:1]
    idx = torch.tensor([stoi[ch] for ch in safe_start], dtype=torch.long, device=device)
    idx = idx.unsqueeze(0)
    h = None
    with torch.no_grad():
        for _ in range(cfg.sample_len):
            logits, h = model(idx, h)
            # 只取最后一个位置的预测
            logits = logits[:, -1, :] / cfg.temperature
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_id], dim=1)
    out = "".join(itos[i] for i in idx[0].tolist())
    return out


def main():
    cfg = Config()
    # 字符表与编码
    chars, stoi, itos = build_vocab(TEXT)
    data = encode(TEXT, stoi)

    # Apple Silicon 上优先用 MPS
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    model = CharRNN(len(chars), cfg.embed_dim, cfg.hidden_dim, cfg.num_layers).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr)

    start_time = time.time()
    model.train()
    for step in range(1, cfg.steps + 1):
        x, y = get_batch(data, cfg)
        x = x.to(device)
        y = y.to(device)

        logits, _ = model(x)
        # 交叉熵：预测下一个字符
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))

        opt.zero_grad()
        loss.backward()
        # 梯度裁剪，防止梯度爆炸
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

        if step % 100 == 0:
            elapsed = time.time() - start_time
            speed = step / max(elapsed, 1e-6)
            print(f"step {step} | loss {loss.item():.4f} | {speed:.1f} steps/s")

        if step % cfg.sample_every == 0:
            print("-" * 60)
            print(sample(model, stoi, itos, device, cfg))
            print("-" * 60)


if __name__ == "__main__":
    # 固定随机种子，便于复现
    random.seed(42)
    torch.manual_seed(42)
    main()