"""
convert_weights.py — 将 Hy-MT2-7B HuggingFace safetensors 转换为 cccc INIReaderBin 格式

用法：
  python convert_weights.py --hf_dir <HF权重目录> --output hy_mt2_7b_fp16_cccc.bin [--dtype half]

依赖：
  pip install safetensors numpy

INIReaderBin 格式说明：
  [0..31]       头部 "CFG_BIN INI"（32字节）
  [32..39]      INI 索引文本长度（uint64 LE）
  [40..40+N-1]  INI 索引文本（每个键对应 "offset,length" 偏移和字节数）
  [40+N..]      二进制数据（各权重依次拼接）
"""

import argparse
import os
import struct
import sys
import numpy as np

try:
    import torch
except ImportError:
    print("请先安装: pip install torch")
    sys.exit(1)

try:
    from safetensors import safe_open
except ImportError:
    print("请先安装: pip install safetensors")
    sys.exit(1)


# ── 权重名称映射：HuggingFace → cccc ─────────────────────────────────────────
# 格式: (hf_name, cccc_name, needs_transpose)
# needs_transpose=False：cccc 使用列主序 BLAS（第一维最快），HF [out,in] C-order
# 在 cccc (in, out, 1, 1) 矩阵中 element(in_i, out_j) 在 in_i + out_j*in_features，
# 恰好与 HF [out,in] C-order 的 out_j*in + in_i 相同（加法交换律），因此不需要转置。
def build_weight_map(num_layers=32):
    mapping = []
    # 全局权重
    mapping.append(("model.embed_tokens.weight", "W_emb", False))  # [V,D] 直接存，embed op 按 id*D 取行
    mapping.append(("lm_head.weight",            "W_lm_head", False))# tied with embed, 保持一致
    mapping.append(("model.norm.weight",          "W_rms_final", False))

    for i in range(num_layers):
        pfx = f"model.layers.{i}"
        mapping.extend([
            (f"{pfx}.input_layernorm.weight",          f"W_rms_attn_{i}", False),
            (f"{pfx}.self_attn.q_proj.weight",         f"W_q_{i}",        False),
            (f"{pfx}.self_attn.k_proj.weight",         f"W_k_{i}",        False),
            (f"{pfx}.self_attn.v_proj.weight",         f"W_v_{i}",        False),
            (f"{pfx}.self_attn.o_proj.weight",         f"W_o_{i}",        False),
            (f"{pfx}.self_attn.query_layernorm.weight",f"W_qnorm_{i}",    False),
            (f"{pfx}.self_attn.key_layernorm.weight",  f"W_knorm_{i}",    False),
            (f"{pfx}.post_attention_layernorm.weight", f"W_rms_ffn_{i}",  False),
            (f"{pfx}.mlp.gate_proj.weight",            f"W_gate_{i}",     False),
            (f"{pfx}.mlp.up_proj.weight",              f"W_up_{i}",       False),
            (f"{pfx}.mlp.down_proj.weight",            f"W_down_{i}",     False),
        ])
    return mapping


# ── INIReaderBin 序列化 ───────────────────────────────────────────────────────
def build_inireaderbin(named_blobs: dict) -> bytes:
    """
    named_blobs: {key: bytes}
    返回完整的 INIReaderBin 二进制内容。
    """
    # 1. 拼接所有二进制 blob，记录 (offset, length)
    blob_parts = []
    index_lines = []
    offset = 0
    for key, data in named_blobs.items():
        blob_parts.append(data)
        index_lines.append(f"{key} = {offset},{len(data)}\n")
        offset += len(data)

    ini_text = "".join(index_lines)
    ini_bytes = ini_text.encode("utf-8")
    binary_content = b"".join(blob_parts)

    # 2. 构建文件
    header = b"CFG_BIN INI" + b"\x00" * (32 - len("CFG_BIN INI"))
    size_ini = struct.pack("<Q", len(ini_bytes))
    return header + size_ini + ini_bytes + binary_content


def convert(hf_dir: str, output_path: str, dtype_str: str = "bfloat16"):
    mapping = build_weight_map(num_layers=32)
    hf_to_cccc = {hf: (cccc, tp) for hf, cccc, tp in mapping}

    # 扫描所有 safetensors 文件
    shard_files = sorted(
        f for f in os.listdir(hf_dir) if f.endswith(".safetensors")
    )
    if not shard_files:
        print(f"在 {hf_dir} 中未找到 .safetensors 文件")
        sys.exit(1)
    print(f"找到 {len(shard_files)} 个 safetensors 分片：{shard_files}")

    # 打开所有分片（lazy load）
    handles = {}
    for sf in shard_files:
        path = os.path.join(hf_dir, sf)
        h = safe_open(path, framework="pt", device="cpu")
        for key in h.keys():
            handles[key] = h

    named_blobs = {}
    # 元信息（data_type 与 cccc ini 中的 data_type 对应）
    named_blobs["data_type"] = dtype_str.encode()
    named_blobs["named_weights"] = b"1"

    missing = []
    for hf_name, cccc_name, needs_transpose in mapping:
        if hf_name not in handles:
            missing.append(hf_name)
            continue

        tensor = handles[hf_name].get_tensor(hf_name)  # torch.Tensor, original dtype (may be bfloat16)

        if needs_transpose and tensor.ndim == 2:
            tensor = tensor.T.contiguous()

        # 转换到目标精度
        if dtype_str == "half":
            tensor = tensor.to(torch.float16)
            blob = tensor.numpy().tobytes()
        elif dtype_str == "bfloat16":
            tensor = tensor.to(torch.bfloat16)
            # numpy 不支持 bfloat16；用 view(int16) 直接取原始字节
            blob = tensor.view(torch.int16).numpy().tobytes()
        elif dtype_str == "float":
            tensor = tensor.to(torch.float32)
            blob = tensor.numpy().tobytes()
        else:
            raise ValueError(f"不支持的 dtype: {dtype_str}")
        named_blobs[f"weight_{cccc_name}"] = blob
        print(f"  {hf_name:60s} → {cccc_name:30s}  shape={tensor.shape}  {len(blob)//1024}KB")

    if missing:
        print(f"\n警告：以下权重未找到：{missing}")

    print(f"\n正在写入 {output_path} ...")
    content = build_inireaderbin(named_blobs)
    with open(output_path, "wb") as f:
        f.write(content)
    print(f"完成！文件大小：{os.path.getsize(output_path) / 1024**3:.2f} GB")


def main():
    parser = argparse.ArgumentParser(description="Hy-MT2-7B → cccc bin 转换工具")
    parser.add_argument("--hf_dir",  required=True, help="HuggingFace 权重目录（含 *.safetensors）")
    parser.add_argument("--output",  default="hy_mt2_7b_bf16_cccc.bin", help="输出 .bin 文件路径")
    parser.add_argument("--dtype",   default="bfloat16", choices=["half", "bfloat16", "float"], help="存储精度（默认 bfloat16，无损）")
    args = parser.parse_args()

    convert(args.hf_dir, args.output, args.dtype)


if __name__ == "__main__":
    main()