"""One-off: InsightFace's ONNX files -> safetensors for nets.py. Needs onnx (only here).

    python convert.py <models dir> <out dir>
"""

import sys
from pathlib import Path

import numpy as np
import onnx
import torch
from onnx import helper, numpy_helper
from safetensors.torch import save_file
from torch import nn

from nets import SCRFD, IResNet

PARAM_OPS = {"Conv": nn.Conv2d, "BatchNormalization": (nn.BatchNorm2d, nn.BatchNorm1d), "PRelu": nn.PReLU, "Gemm": nn.Linear}


def attrs(n):
    return {a.name: helper.get_attribute_value(a) for a in n.attribute}


def load(model: nn.Module, path: Path) -> nn.Module:
    g = onnx.load(str(path)).graph
    init = {t.name: torch.from_numpy(numpy_helper.to_array(t).copy()) for t in g.initializer}
    mods = [m for m in model.modules() if isinstance(m, (nn.Conv2d, nn.BatchNorm2d, nn.BatchNorm1d, nn.PReLU, nn.Linear))]
    scales = [m for m in model.modules() if hasattr(m, "scale") and isinstance(getattr(m, "scale"), torch.Tensor)]
    nodes = [n for n in g.node if n.op_type in PARAM_OPS]
    assert len(nodes) == len(mods), (len(nodes), len(mods))
    for n, m in zip(nodes, mods):
        assert isinstance(m, PARAM_OPS[n.op_type]), (n.name, n.op_type, m)
        w = [init[i] for i in n.input[1:]]
        a = attrs(n)
        if n.op_type == "Conv":
            assert a.get("group", 1) == 1 and a.get("dilations", [1, 1]) == [1, 1]
            assert list(a["strides"]) == list(m.stride) and list(a["pads"]) == list(m.padding) * 2, (n.name, a, m)
            m.weight.data.copy_(w[0]), m.bias.data.copy_(w[1])
        elif n.op_type == "BatchNormalization":
            assert abs(a["epsilon"] - m.eps) < 1e-9
            for t, v in zip((m.weight, m.bias, m.running_mean, m.running_var), w):
                t.data.copy_(v)
        elif n.op_type == "PRelu":
            m.weight.data.copy_(w[0].reshape(-1))
        else:
            assert a.get("transB") == 1 and a.get("alpha", 1.0) == 1.0 and a.get("beta", 1.0) == 1.0
            m.weight.data.copy_(w[0]), m.bias.data.copy_(w[1])
    muls = [n for n in g.node if n.op_type == "Mul"]
    assert len(muls) == len(scales)
    for n, s in zip(muls, scales):
        s.scale.copy_(init[n.input[1]])
    # what the module hard-codes: nearest 2x upsampling, the heads' reshapes
    for n in g.node:
        if n.op_type == "Resize":
            print("  resize", {k: (v.decode() if isinstance(v, bytes) else v) for k, v in attrs(n).items()})
        if n.op_type == "Reshape":
            assert init[n.input[1]].tolist() in ([-1, 1], [-1, 4], [-1, 10]), init[n.input[1]]
    if muls:
        print("  head scales", [float(s.scale) for s in scales])
    return model.eval()


def main(src: Path, out: Path):
    out.mkdir(parents=True, exist_ok=True)
    for name, model, file in [("det_10g", SCRFD(), "buffalo_l/det_10g.onnx"),
                              ("w600k_r50", IResNet(IResNet.LAYERS["r50"]), "buffalo_l/w600k_r50.onnx"),
                              ("glintr100", IResNet(IResNet.LAYERS["r100"]), "antelopev2/antelopev2/glintr100.onnx")]:
        print(name)
        load(model, src / file)
        # the recognition models in fp16 (half the download, embeddings unchanged to 1e-6); the small
        # detector in fp32, whose boxes would move by a tenth of a pixel
        dtype = torch.float32 if isinstance(model, SCRFD) else torch.float16
        save_file({k: v.to(dtype).contiguous() for k, v in model.state_dict().items() if not k.endswith("num_batches_tracked")},
                  str(out / f"{name}.safetensors"), metadata={"source": file, "converted_by": "ZenTrainer convert.py"})


if __name__ == "__main__":
    main(Path(sys.argv[1]), Path(sys.argv[2]))