#!/usr/bin/env python3
"""训练对话模型——清洗语料 + 可恢复分阶段训练

策略：
  1. 递归读取 corpus/ 下的本地语料
  2. 清理 markdown/HTML 噪声，过滤过短或过长段落
  3. 同时投喂简体和繁体版本
  4. 统一使用 LanguageTrainingSession 训练、验证和保存
  5. 分阶段降低学习率
"""
import torch
from pathlib import Path
import time
import re
import os

from jspaceai import (
    LanguageConfig, JSpaceLanguageModel,
    CharTokenizer, load_chinese_corpus,
    LanguageTrainingConfig, LanguageTrainingSession,
)


def get_config(vocab_size: int) -> LanguageConfig:
    return LanguageConfig(
        vocab_size=vocab_size, embed_dim=64, input_dim=32,
        workspace_dim=128, expert_dim=64, num_experts=12, num_wells=8,
        ode_steps=4, dt=0.1, tau_w=0.6, jacobian_sparsity=32, noise_std=0.001,
        use_rk4=True, use_layer_norm=True,
    )


def _corpus_budget_mb(default: int | None = 64) -> int | None:
    raw = os.environ.get("JSPACE_CORPUS_MAX_MB")
    if raw is None:
        return default
    raw = raw.strip().lower()
    if raw in {"", "0", "all", "none"}:
        return None
    return int(raw)


def clean_corpus(max_source_mb: int | None = None, use_cache: bool = True) -> str:
    """构建中文训练语料——读取 corpus/ 下所有文件，繁简双版本同时投喂。

    策略：
      1. 递归读取 corpus/ 目录下的本地语料（默认有读取预算）
      2. 去掉 markdown/HTML 标记（链接、表格、标题符号等），只保留纯文本
      3. 对每段文本同时生成简体版和繁体版，都喂给模型
      4. 加上内嵌唐诗宋词论语（简体连续文本）
    """
    from opencc import OpenCC

    if max_source_mb is None:
        max_source_mb = _corpus_budget_mb()

    cache_key = "all" if max_source_mb is None else f"{max_source_mb}mb"
    cache_path = Path("outputs/cache") / f"clean_corpus_{cache_key}.txt"
    if use_cache and cache_path.exists():
        return cache_path.read_text(encoding="utf-8")

    cc_s2t = OpenCC('s2t')  # 简转繁
    cc_t2s = OpenCC('t2s')  # 繁转简

    corpus_dir = Path(__file__).parent / 'corpus'
    raw_paragraphs = []
    source_budget = None if max_source_mb is None else max_source_mb * 1024 * 1024
    bytes_read = 0

    # 1. 内嵌唐诗宋词论语
    for para in load_chinese_corpus().split('\n\n'):
        para = para.strip()
        if para:
            raw_paragraphs.append(para)

    # 2. LCCC 对话语料（最高优先级，真实对话）
    lccc_file = corpus_dir / 'lccc_dialogues.txt'
    if lccc_file.exists():
        lccc_text = lccc_file.read_text(encoding='utf-8')
        for para in lccc_text.split('\n\n'):
            para = para.strip()
            if para:
                raw_paragraphs.append(para)

    # 3. 递归读取 corpus/ 下其他所有文件（教材+百科）
    if corpus_dir.exists():
        for filepath in sorted(corpus_dir.rglob('*')):
            if not filepath.is_file():
                continue
            if source_budget is not None:
                try:
                    file_size = filepath.stat().st_size
                except OSError:
                    continue
                if bytes_read >= source_budget:
                    break
                if file_size > max(1024 * 1024, source_budget - bytes_read):
                    continue
            try:
                content = filepath.read_text(encoding='utf-8', errors='ignore')
            except Exception:
                continue
            bytes_read += len(content.encode('utf-8', errors='ignore'))
            # 去 markdown 标记
            content = re.sub(r'\[([^\]]*)\]\([^)]*\)', r'\1', content)  # 链接
            content = re.sub(r'^#{1,6}\s+', '', content, flags=re.M)    # 标题
            content = re.sub(r'^\|.*\|$', '', content, flags=re.M)      # 表格行
            content = re.sub(r'^---+$', '', content, flags=re.M)        # 分隔线
            content = re.sub(r'```[^`]*```', '', content, flags=re.S)   # 代码块
            content = re.sub(r'\*\*([^*]+)\*\*', r'\1', content)        # 粗体
            content = re.sub(r'\*([^*]+)\*', r'\1', content)            # 斜体
            content = re.sub(r'^- ', '', content, flags=re.M)           # 列表
            # 按空行分段
            for para in content.split('\n\n'):
                para = para.strip()
                if para:
                    raw_paragraphs.append(para)

    # 清洗过滤
    cleaned = []
    for para in raw_paragraphs:
        para = re.sub(r'[ \t]+', ' ', para)
        para = re.sub(r'\n{3,}', '\n\n', para)
        para = para.strip()
        if len(para) < 10 or len(para) > 500:
            continue
        # 至少要有一些中文字符
        zh_chars = sum(1 for c in para if len(c) == 1 and ord(c) > 0x4e00)
        if zh_chars < 3:
            continue
        cleaned.append(para)

    # 繁简双版本投喂
    bilingual = []
    for para in cleaned:
        simplified = cc_t2s.convert(para)
        traditional = cc_s2t.convert(simplified)
        bilingual.append(simplified)
        bilingual.append(traditional)

    result = '\n\n'.join(bilingual)
    if use_cache:
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        cache_path.write_text(result, encoding="utf-8")
    return result


def main():
    text = clean_corpus()
    print(f"清洗后语料: {len(text)} 字符")

    tok = CharTokenizer.from_text(text, max_vocab=3000)  # 覆盖98.9%文本
    cfg = get_config(tok.vocab_size)
    model = JSpaceLanguageModel(cfg)
    print(f"vocab={tok.vocab_size}, params={sum(p.numel() for p in model.parameters()):,}")

    all_tokens = tok.encode(text)

    # 加载已有模型
    mp = Path('outputs/chat_model.pt')
    if mp.exists():
        try:
            ckpt = torch.load(mp, map_location='cpu', weights_only=False)
            if ckpt['config'].vocab_size == cfg.vocab_size:
                model.load_state_dict(ckpt['model'])
                print(f"已加载: {mp}")
        except Exception:
            print("加载失败，全新训练")

    prompts = ['学而时习之', '床前明月光', '學而時習之', '春天', '人']

    stages = [
        (3e-3, 500, "快速下降"),
        (1e-3, 500, "稳定收敛"),
        (5e-4, 500, "精调"),
        (2e-4, 500, "最终精调"),
    ]

    t_total = time.time()
    for lr, n_steps, label in stages:
        print(f"\n{'='*60}")
        print(f"阶段: {label} | lr={lr} | {n_steps}步")
        print(f"{'='*60}")
        train_cfg = LanguageTrainingConfig(
            seq_len=64,
            batch_size=8,
            lr=lr,
            ewc_lambda=0.05,
            consolidate_every=100,
            validate_every=100,
            save_every=100,
            use_euler_during_train=True,
        )
        trainer = LanguageTrainingSession(
            model, cfg, tok, train_cfg, device='cpu',
        )

        def report(stats: dict):
            if stats['step'] % 100 != 0:
                return
            val = f" val={stats['val_loss']:.3f}" if "val_loss" in stats else ""
            print(f"\n  step {stats['step']:3d} | loss {stats['loss']:.3f}{val}")
            model.eval()
            for prompt in prompts:
                generated = model.generate(tok.encode(prompt) or [0], n_new=40, temperature=0.7, top_k=5)
                print(f"    [{prompt}] {repr((prompt + tok.decode(generated))[:60])}")
            model.train()

        trainer.fit_tokens(
            all_tokens,
            max_steps=n_steps,
            checkpoint_path=mp,
            on_progress=report,
        )

    print(f"\n完成，总耗时 {time.time()-t_total:.0f}s，已保存到 outputs/chat_model.pt")


if __name__ == '__main__':
    main()
