#!/usr/bin/env python3
"""
inline_onnx.py
==============

把 `.model4` 里的 ONNX 图**内联成独立可跑的 .onnx**（不依赖 iZotope 运行时）。

背景
----
ONNX 图里大部分权重是 `external_data`，`location` 写成 `model4_<hex(张量名)>.bin`——
这些文件在磁盘上根本不存在，是 iZotope 的 Photon 在运行时从 safetensors 载荷里喂进去的。
所以抽出来的 .onnx 用 onnxruntime 打不开（`should be stored in model4_xxx.bin, but it is not regular file`）。

本脚本用容器自带的 `graphs[backend=onnx].bindings.tensors` 表把
`backend_name -> 容器张量名` 解析出来，把字节写回 initializer，输出标准 .onnx。

张量来源（三级回退）
--------------------
1. 容器里同名张量（规范名）
2. `bindings.tensors` 里的 `backend_name` 反查
3. `onnx.source.<sha>.<name>` 去掉前缀后的同名张量

用法
----
    python inline_onnx.py                 # 处理非 pymss 的 10 个模型
    python inline_onnx.py --all           # 连 pymss_models 里的 4 个也一起导
    python inline_onnx.py --only indoor
"""
from __future__ import annotations

import argparse
import json
import os
import re
import struct
import sys

import numpy as np
import onnx
import torch

import model_io
from onnx import numpy_helper

HERE = os.path.dirname(os.path.abspath(__file__))
SF_DIR = os.path.join(HERE, "..", "raw", "safetensors")
BLOBS_DIR = os.path.join(HERE, "..", "raw", "blobs")
CKPT_DIR = os.path.join(HERE, "..", "raw", "ckpt")
OUT_DIR = os.path.join(HERE, "..", "onnx")

# 已经塞进 pymss（stock BSRoformer 直接 strict 可加载）的 4 个
PYMSS_OK = {
    "voice2stem_v03_fp16c_t2", "voicesep2stem_fp16c_t2",
    "drumsep_1.3", "music4stemhq_fp16c_t2",
}


def read_header(model: str) -> dict:
    p = model_io.safetensors_path(model)
    with open(p, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        hdr = json.loads(f.read(n))
    return json.loads(hdr["__metadata__"]["model4"])


def container_tensors(model: str) -> dict:
    sd = model_io.state_dict_for(model)
    out = {}
    for k, v in sd.items():
        if k.startswith("__blob__"):
            continue
        out[k] = v
        # 去掉 onnx.source.<sha>. / coreml.* 前缀后也登记一份
        m = re.match(r"^(?:onnx|coreml)\.source\.[0-9a-f]+\.(.*)$", k)
        if m:
            out.setdefault(m.group(1), v)
        m = re.match(r"^coreml\.fp16\.(.*)_to_fp16$", k)
        if m:
            out.setdefault(m.group(1), v)
    return out


def norm_name(k: str) -> str:
    n = k.replace("/", ".").replace("..", ".").strip(".")
    n = re.sub(r"^model\.", "", n)
    n = re.sub(r"^__onnx__\.", "", n)
    return n


def build_lookup(tensors: dict) -> dict:
    """归一化名 -> tensor，供名字只差分隔符/前缀时兜底。"""
    out = {}
    for k, v in tensors.items():
        out.setdefault(norm_name(k), v)
    return out


def binding_map(meta: dict) -> dict:
    """backend_name -> 容器张量名"""
    g = next((x for x in meta.get("graphs", []) if x.get("backend") == "onnx"), None)
    if not g:
        return {}
    out = {}
    for b in g.get("bindings") or []:
        for tname, v in (b.get("tensors") or {}).items():
            bn = v.get("backend_name")
            if bn and bn not in out:
                out[bn] = tname
    return out


def pick_onnx_blob(model: str):
    bd = model_io.blobs_dir(model)
    if not os.path.isdir(bd):
        return None
    cands = sorted((p for p in os.listdir(bd) if p.endswith(".onnx")),
                   key=lambda n: os.path.getsize(os.path.join(bd, n)), reverse=True)
    for n in cands:
        p = os.path.join(bd, n)
        try:
            m = onnx.load(p, load_external_data=False)
        except Exception:
            continue
        if len(m.graph.initializer) > 0:
            return p, m
    return None, None


def inline(model: str, tensors: dict, meta: dict):
    path, m = pick_onnx_blob(model)
    if m is None:
        return None, {"error": "没有可解析的 onnx blob"}

    bmap = binding_map(meta)
    look = build_lookup(tensors)
    # roformer 系（voice_restore 等）的 ONNX 权重存在规范名 blk.*/band_split.* 下，
    # 借用 to_pymss 的映射把 model.layers.* 还原回容器名
    try:
        sys.path.insert(0, HERE)
        from to_pymss import load_canon, map_bs_roformer
        ck = model_io.model_path_for(model)
        if ck and os.path.exists(ck):
            canon, al = load_canon(ck)
            if any(k.startswith("band_split") for k in canon):
                sd, _p, _u = map_bs_roformer(canon, al)
                for k, v in sd.items():
                    look.setdefault(norm_name("model." + k), v)
                    # 部分模型把注意力包在 .branch. 里
                    m2 = re.match(r"^(layers\.\d+\.\d+\.layers\.\d+\.\d+)\.(.+)$", k)
                    if m2:
                        look.setdefault(norm_name(f"model.{m2.group(1)}.branch.{m2.group(2)}"), v)
    except Exception as e:
        print("    (roformer 还原跳过:", e, ")")
    aliases = meta.get("tensor_aliases") or {}

    def deref(n):
        """沿 tensor_aliases 追到实际存储的张量名。"""
        seen = set()
        while isinstance(n, str) and n in aliases and n not in seen:
            seen.add(n)
            n = (aliases[n] or {}).get("tensor", n)
        return n
    need = [i for i in m.graph.initializer if any(e.key == "location" for e in i.external_data)]
    resolved = missing = 0
    for init in m.graph.initializer:
        loc = next((e.value for e in init.external_data if e.key == "location"), None)
        if loc is None:
            continue
        name = init.name
        t = tensors.get(name)
        if t is None and name in bmap:
            t = tensors.get(bmap[name])
        if t is None:
            try:
                dec = bytes.fromhex(loc.split("model4_")[-1].replace(".bin", "")).decode()
                t = tensors.get(dec) or tensors.get(bmap.get(dec, ""))
            except Exception:
                t = None
        if t is None:
            # Bandit：ONNX 里是 combined.0(LayerNorm) / combined.1(Linear)，
            # 容器里是分开的 norm.* / fc.*
            m2 = re.match(r"^model\.(band_split\.norm_fc_modules\.\d+)\.combined\.([01])\.(\w+)$", name)
            if m2:
                sub = "norm" if m2.group(2) == "0" else "fc"
                t = look.get(norm_name(f"{m2.group(1)}.{sub}.{m2.group(3)}"))
        if t is None:
            t = look.get(norm_name(deref(name)))
        if t is None:
            t = look.get(norm_name(name))
        if t is None and name in bmap:
            # 关键一步：binding 的目标名往往还只是一层别名
            #   onnx.source.<sha>.model.band_split...  ->  band_split.0.norm.weight
            t = look.get(norm_name(deref(bmap[name])))
        if t is None and name in bmap:
            t = look.get(norm_name(bmap[name]))
        if t is None:
            missing += 1
            continue
        want = onnx.helper.tensor_dtype_to_np_dtype(init.data_type)
        arr = t.detach().cpu().numpy()
        if arr.dtype != want:
            arr = arr.astype(want, copy=False)
        if tuple(init.dims) and tuple(arr.shape) != tuple(init.dims):
            if arr.size != int(np.prod(init.dims)):
                missing += 1
                continue
            arr = arr.reshape(init.dims)
        init.CopyFrom(numpy_helper.from_array(arr, name=name))
        resolved += 1

    # 第二遍：首轮因解析顺序漏掉的再扫一次（按归一化名直接取）
    for init in m.graph.initializer:
        if not any(e.key == "location" for e in init.external_data):
            continue
        t2 = look.get(norm_name(init.name))
        if t2 is None:
            continue
        want2 = onnx.helper.tensor_dtype_to_np_dtype(init.data_type)
        arr2 = t2.detach().cpu().numpy()
        if arr2.dtype != want2:
            arr2 = arr2.astype(want2, copy=False)
        if tuple(init.dims) and tuple(arr2.shape) != tuple(init.dims):
            if arr2.size != int(np.prod(init.dims)):
                continue
            arr2 = arr2.reshape(init.dims)
        init.CopyFrom(numpy_helper.from_array(arr2, name=init.name))
        resolved += 1
        missing = max(0, missing - 1)

    os.makedirs(OUT_DIR, exist_ok=True)
    dst = os.path.join(OUT_DIR, model + ".onnx")   # 平铺，不套文件夹
    onnx.save(m, dst)
    info = {"model": model, "source_blob": os.path.basename(path),
            "nodes": len(m.graph.node), "initializers": len(m.graph.initializer),
            "external_refs": len(need), "resolved": resolved, "missing": missing,
            "bytes": os.path.getsize(dst)}
    # 校验
    try:
        onnx.checker.check_model(dst)
        info["checker"] = "OK"
    except Exception as e:
        info["checker"] = str(e)[:160]
    return dst, info


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--all", action="store_true", help="包含已塞进 pymss 的 4 个")
    ap.add_argument("--only", default=None)
    args = ap.parse_args()
    os.makedirs(OUT_DIR, exist_ok=True)

    models = model_io.list_models()
    if not args.all:
        models = [n for n in models if n not in PYMSS_OK]
    if args.only:
        models = [n for n in models if args.only.lower() in n.lower()]

    report = []
    for model in models:
        try:
            meta = read_header(model)
            tensors = container_tensors(model)
            dst, info = inline(model, tensors, meta)
        except Exception as e:
            info = {"model": model, "error": f"{type(e).__name__}: {e}"}
            dst = None
        tag = "OK" if info.get("missing") == 0 and info.get("checker") == "OK" else "!!"
        print(f"[{tag}] {model:<32} " + (
            f"nodes={info.get('nodes')} inits={info.get('initializers')} "
            f"外部={info.get('external_refs')} 解析={info.get('resolved')} "
            f"缺={info.get('missing')} checker={info.get('checker')}"
            if "error" not in info else info["error"]))
        report.append(info)

    with open(os.path.join(OUT_DIR, "_inline_report.json"), "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2, ensure_ascii=False)
    ok = sum(1 for r in report if r.get("missing") == 0 and r.get("checker") == "OK")
    print(f"\n完整内联 {ok}/{len(report)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())