"""把 DomSense 训练 checkpoint 转换为 Hugging Face 标准格式并导出。

用法:
    python scripts/convert_to_hf.py
        --ckpt ../outputs/rlcd_mdp_run1/best_model.pt
        --out  ./saved_model

产出（HF 标准目录）:
    saved_model/
        config.json          # SystemOneConfig 序列化
        model.safetensors    # 转换后的权重（自动去掉 base_model_prefix "core."）
        tokenizer.json       # DomSenseTokenizer

导出后，可这样加载:
    import sys; sys.path.insert(0, "repo_root")
    from huggingface_repo import SystemOneModelForDecision, DomSenseTokenizer
    m = SystemOneModelForDecision.from_pretrained("saved_model_dir")
"""

import argparse
import sys
from pathlib import Path

import torch

REPO_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(REPO_ROOT))

from huggingface_repo import SystemOneModelForDecision, DomSenseTokenizer
from huggingface_repo.configuration_domsense import SystemOneConfig


def convert(ckpt_path: str, out_dir: str):
    print(f"[1/4] 读取 checkpoint: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location="cpu")
    state_dict = ckpt["model_state_dict"]

    # 兼容两种 checkpoint：扁平 dict 或 {model: {...}} 嵌套
    raw_cfg = ckpt.get("config", {})
    if isinstance(raw_cfg, dict) and "model" in raw_cfg:
        raw_cfg = raw_cfg["model"]
    cfg_kwargs = {
        k: v
        for k, v in raw_cfg.items()
        if v is not None and not isinstance(v, dict)
    }
    # 用权重形状自动修正关键结构维度，避免 config 与 checkpoint 不一致
    if "mdp_head.action_embed_table.weight" in state_dict:
        cfg_kwargs["max_actions"] = state_dict[
            "mdp_head.action_embed_table.weight"
        ].shape[0]
    # 领域数量由权重推断；领域名优先从 config 保留（若有），否则自动命名
    if "domain_moe.gate.2.weight" in state_dict:
        n_slots = state_dict["domain_moe.gate.2.weight"].shape[0]
        n_domains = n_slots - 1
        cfg_kwargs["use_domain_experts"] = n_domains > 0
        if "domains" not in cfg_kwargs or not cfg_kwargs["domains"]:
            cfg_kwargs["domains"] = [f"domain_{i}" for i in range(n_domains)]
        # 确保领域列表长度与权重匹配
        if len(cfg_kwargs["domains"]) != n_domains:
            cfg_kwargs["domains"] = [f"domain_{i}" for i in range(n_domains)]
    cfg = SystemOneConfig(**cfg_kwargs)

    print(f"[2/4] 构建 HF 模型；hidden_dim={cfg.hidden_dim}, "
          f"num_bellman_steps={cfg.num_bellman_steps}, "
          f"use_domain_experts={cfg.use_domain_experts}")
    model = SystemOneModelForDecision(cfg)

    print("[3/4] 加载权重到内核 core.*（strict=False）")
    missing, unexpected = model.core.load_state_dict(state_dict, strict=False)
    if missing:
        print(f"  - missing {len(missing)} keys: {missing[:10]}")
    if unexpected:
        print(f"  - unexpected {len(unexpected)} keys: {unexpected[:10]}")
    total, trainable = model.count_parameters()
    print(f"  总参数量: {total:,} | 可训练: {trainable:,}")

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    print(f"[4/4] 导出 HF 格式到: {out}")
    # save_pretrained 基于 base_model_prefix 自动存储（保存去 core.，加载加回）
    model.save_pretrained(str(out))
    DomSenseTokenizer().save_pretrained(str(out))

    print("完成！现在可以从该目录加载:")
    print(f'    from huggingface_repo import SystemOneModelForDecision, DomSenseTokenizer')
    print(f'    m = SystemOneModelForDecision.from_pretrained("{out}")')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", default="../outputs/rlcd_mdp_run1/best_model.pt")
    parser.add_argument("--out", default="./saved_model")
    args = parser.parse_args()
    convert(args.ckpt, args.out)