#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
用已有的 CharTokenizer + build_vocabulary 从训练数据生成词表
然后直接启动训练
"""
import sys, os, json, time
sys.path.insert(0, '/workspace/xiaoli-llm-chat')
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', buffering=1)

from llm_chat_improved import tokenizer_instance, build_vocabulary

DATA_PATH = '/workspace/data/merged_clean_v3.json'
VOCAB_OUT = '/workspace/vocab.json'
VOCAB_OUT_XIAOLI = '/workspace/xiaoli-llm-chat/vocab.json'

def load_pairs(path):
    """加载训练数据为 (question, answer) 对"""
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    pairs = []
    for item in data:
        if isinstance(item, dict):
            q = item.get("question", "") or item.get("input", "") or item.get("prompt", "")
            a = item.get("answer", "") or item.get("output", "") or item.get("response", "")
            if q and a:
                pairs.append((q, a))
    return pairs

def main():
    print("=" * 60)
    print("用 build_vocabulary 从训练数据生成词表")
    print("=" * 60)
    
    # 1. 加载训练数据
    t0 = time.time()
    print("\n[1] 加载训练数据...")
    pairs = load_pairs(DATA_PATH)
    print(f"  数据量: {len(pairs)} 对")
    print(f"  耗时: {time.time()-t0:.1f}s")
    
    # 2. 用 build_vocabulary 生成词表
    t1 = time.time()
    print("\n[2] 用 build_vocabulary 分析数据生成词表...")
    vocab = build_vocabulary(pairs, min_freq=1)
    
    # build_vocabulary 返回的是 tokenizer_instance.char_to_id
    # 已包含特殊token (<PAD>=0, <UNK>=1, etc.)
    special_tokens = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
                      '<START>': 4, '<END>': 5, '<SEP>': 6}
    
    real_vocab_size = max(vocab.values()) + 1
    
    # 统计
    single_char_tokens = sum(1 for k in vocab if len(k) == 1)
    multi_char_tokens = sum(1 for k in vocab if len(k) > 1 and not k.startswith('<'))
    special_token_count = sum(1 for k in vocab if k.startswith('<'))
    
    print(f"  词表总项: {len(vocab)}")
    print(f"  vocab_size(模型用): {real_vocab_size}")
    print(f"  单字符token: {single_char_tokens}")
    print(f"  多字符token(英文词): {multi_char_tokens}")
    print(f"  特殊token: {special_token_count}")
    print(f"  耗时: {time.time()-t1:.1f}s")
    
    # 3. 编码测试
    print("\n[3] 编码测试:")
    test_texts = [
        "你好呀！我是小狸～",
        "小狸最喜欢吃草莓冰淇淋了！",
        "<SOS>你好<SEP><START>我是小狸～<END><EOS>",
    ]
    for text in test_texts:
        ids = tokenizer_instance.encode(text)
        decoded = tokenizer_instance.decode(ids)
        print(f"  原文: '{text}'")
        print(f"  编码: {ids[:20]}... (长度={len(ids)})")
        print(f"  解码: '{decoded}'")
        print()
    
    # 4. 估算模型参数量
    print("\n[4] 模型参数量估算 (d_model=320, n_layers=6, d_ff=1536):")
    d_model = 320
    n_layers = 6
    d_ff = 1536
    vs = real_vocab_size
    emb = vs * d_model
    out = vs * d_model
    per_layer = 4 * d_model**2 + 3 * d_model * d_ff + 2 * d_model
    total = emb + out + n_layers * per_layer + 2 * d_model
    print(f"  vocab_size={vs}")
    print(f"  Embedding: {emb:,}")
    print(f"  fc_out: {out:,}")
    print(f"  Transformer({n_layers}层): {n_layers*per_layer:,}")
    print(f"  总参数量: {total:,} ({total/1e6:.2f}M)")
    
    # 5. 保存词表
    print(f"\n[5] 保存词表...")
    
    # vocab.json 格式：不含特殊token，id从7开始（与 train.py 兼容）
    normal_vocab = {k: v for k, v in vocab.items() if k not in special_tokens}
    
    # 备份旧词表
    for path in [VOCAB_OUT, VOCAB_OUT_XIAOLI]:
        if os.path.exists(path):
            backup = path.replace('vocab.json', 'vocab_old.json')
            import shutil
            shutil.copy2(path, backup)
            print(f"  旧词表已备份: {backup}")
    
    # 保存到两个位置
    for path in [VOCAB_OUT, VOCAB_OUT_XIAOLI]:
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(normal_vocab, f, ensure_ascii=False, indent=1)
        print(f"  {path}: {len(normal_vocab)}项 (不含特殊token, id从7开始)")
    
    # 保存完整词表（含特殊token，用于参考）
    full_path = VOCAB_OUT_XIAOLI.replace('vocab.json', 'vocab_full.json')
    with open(full_path, 'w', encoding='utf-8') as f:
        json.dump(vocab, f, ensure_ascii=False, indent=1)
    print(f"  {full_path}: {len(vocab)}项 (含特殊token)")
    
    total_time = time.time() - t0
    print(f"\n总耗时: {total_time:.1f}s")
    print("=" * 60)
    print(f"词表生成完成！vocab_size = {real_vocab_size}")
    print(f"模型参数量 ≈ {total/1e6:.2f}M")
    print("=" * 60)


if __name__ == "__main__":
    main()