#!/usr/bin/env python3
"""Export NadirNet to ONNX: the page graph (palette + page context) and the element graph (treatment + role for any
set of page images). Calibrated 8-bit and 4-bit variants of the page graph come from quantize.py."""
import argparse, json, os, pickle, time
import numpy as np
import torch
import torch.nn as nn

from model import NadirNet, ElemNet, FI, TREAT, IMG_ROLES, N_CTX
from train import collate, prep, element_logits

NAMES_IN = ["tok_f", "tok_role", "tok_lab", "img_f", "img_thumb", "page_f", "src", "rel", "mask"]
NAMES_OUT = ["lab", "alpha_logit", "ctx"]
ELEM_IN = ["e_f", "e_thumb", "under_light", "under_dark", "ctx"]
ELEM_OUT = ["treat_logits", "role_logits"]


class Wrapped(nn.Module):
    def __init__(self, m):
        super().__init__()
        self.m = m

    def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
        return self.m(tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask)


def feeds_for(s, pad=None):
    b = collate([s], "cpu")
    if pad:  # fixed shapes for calibration (quantizers stack activations across samples)
        N0, M0 = b["tok_f"].shape[1], b["img_f"].shape[1]
        N, M = pad
        if N0 <= N and M0 <= M:
            T0, T = 1 + N0 + M0, 1 + N + M
            def padto(x, dims, val=0):
                shape = list(x.shape)
                for d, sz in dims.items():
                    shape[d] = sz
                y = torch.full(shape, val, dtype=x.dtype)
                y[tuple(slice(0, x.shape[i]) for i in range(x.ndim))] = x
                return y
            idx = list(range(1 + N0)) + list(range(1 + N, 1 + N + M0))
            rel = torch.zeros(1, T, T, b["rel"].shape[-1]); mask = torch.zeros(1, T, dtype=torch.bool)
            ii = torch.tensor(idx)
            rel[0][ii[:, None], ii[None, :]] = b["rel"][0]
            mask[0, ii] = b["mask"][0]
            b = dict(b, tok_f=padto(b["tok_f"], {1: N}), role=padto(b["role"], {1: N}), lab=padto(b["lab"], {1: N}),
                     img_f=padto(b["img_f"], {1: M}), thumbs=padto(b["thumbs"], {1: M}), rel=rel, mask=mask)
    M = b["img_f"].shape[1]
    if M == 0:  # the runtime always sends at least one (masked) image slot
        T = b["mask"].shape[1]
        b["img_f"] = torch.zeros(1, 1, FI); b["thumbs"] = torch.zeros(1, 1, 4, 32, 32)
        rel = torch.zeros(1, T + 1, T + 1, b["rel"].shape[-1]); rel[:, :T, :T] = b["rel"]; b["rel"] = rel
        b["mask"] = torch.cat([b["mask"], torch.zeros(1, 1, dtype=torch.bool)], 1)
    return {"tok_f": b["tok_f"].numpy(), "tok_role": b["role"].numpy(), "tok_lab": b["lab"].numpy(), "img_f": b["img_f"].numpy(),
            "img_thumb": b["thumbs"].numpy(), "page_f": b["page"].numpy(), "src": b["src"].numpy(), "rel": b["rel"].numpy(), "mask": b["mask"].numpy()}


def elem_feeds(s, lab, ctx):
    """Element-graph inputs for a page's labelled images, with the dark surface taken from the page graph's output."""
    b = collate([s], "cpu")
    k = int(b["e_mask"][0].sum())
    ut = b["e_ut"][0, :k].numpy()
    return {"e_f": b["e_f"][:, :k].numpy(), "e_thumb": b["e_thumb"][:, :k].numpy(), "under_light": b["e_ul"][:, :k].numpy(),
            "under_dark": lab[:, ut].astype(np.float32), "ctx": ctx.astype(np.float32)}, b["e_soft"][0, :k].numpy()


def fidelity(sess, ref_sess, samples):
    """Mean OKLab distance between this model's colours and fp32's (weighted like training), worst token, ctx drift."""
    de, worst, ctxd = [], 0.0, []
    for s in samples:
        f = feeds_for(s)
        a = sess.run(None, f); r = ref_sess.run(None, f)
        n = s["F"].shape[0]
        d = np.linalg.norm(a[0][0, :n] - r[0][0, :n], axis=-1)
        w = s["w"][:n]
        de.append(float((d * w).sum() / w.sum())); worst = max(worst, float(d.max()))
        ctxd.append(float(np.abs(a[2] - r[2]).max()))
    return {"dE_vs_fp32": float(np.mean(de)), "dE_max": worst, "ctx_max_abs": float(np.max(ctxd))}


def bench(sess, feeds, reps=3):
    ts = []
    for f in feeds[:20]:
        sess.run(None, f)
        t = time.time()
        for _ in range(reps):
            sess.run(None, f)
        ts.append((time.time() - t) / reps * 1000)
    return float(np.median(ts))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--only-float", action="store_true", help="kept for the job scripts (quantization lives in quantize.py)")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    import onnx, onnxruntime as ort
    ck = torch.load(args.ckpt, map_location="cpu")
    model = NadirNet(**{k: v for k, v in ck["cfg"].items()}); model.load_state_dict(ck["model"]); model.eval()
    data = pickle.load(open(args.data, "rb"))
    val = [prep(s) for s in data["val"]]
    s0 = max(val, key=lambda s: len(s["role"]))
    ex = feeds_for(s0)
    par_samples = val[:8]
    with torch.no_grad():
        ref_out = [model(*[torch.from_numpy(feeds_for(s)[k]) for k in NAMES_IN]) for s in par_samples]
    fp32 = os.path.join(args.out, "nadir-fp32.onnx")
    torch.onnx.export(Wrapped(model), tuple(torch.from_numpy(ex[k]) for k in NAMES_IN), fp32, input_names=NAMES_IN, output_names=NAMES_OUT,
                      dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"},
                                    "rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}},
                      opset_version=17, do_constant_folding=True, dynamo=False)
    onnx.checker.check_model(fp32)
    so = ort.SessionOptions(); so.intra_op_num_threads = 4
    ref = ort.InferenceSession(fp32, so, providers=["CPUExecutionProvider"])
    parity = 0.0
    for r, s in zip(ref_out, par_samples):
        o = ref.run(None, feeds_for(s))
        parity = max(parity, float(np.abs(r[0].numpy() - o[0]).max()), float(np.abs(r[2].numpy() - o[2]).max()))
    # element graph
    es = [s for s in val if len(s["E_ut"])][:64]
    ef = []
    for s in es:
        o = ref.run(None, feeds_for(s))
        ef.append(elem_feeds(s, o[0], o[2]))
    exe = max(ef, key=lambda x: x[0]["e_f"].shape[1])[0] if ef else {"e_f": np.zeros((1, 4, FI), np.float32), "e_thumb": np.zeros((1, 4, 4, 32, 32), np.float32),
                                                                    "under_light": np.zeros((1, 4, 3), np.float32), "under_dark": np.zeros((1, 4, 3), np.float32), "ctx": np.zeros((1, N_CTX), np.float32)}
    elem = os.path.join(args.out, "nadir-elem-fp32.onnx")
    em = ElemNet(model).eval()
    torch.onnx.export(em, tuple(torch.from_numpy(exe[k]) for k in ELEM_IN), elem, input_names=ELEM_IN, output_names=ELEM_OUT,
                      dynamic_axes={"e_f": {1: "K"}, "e_thumb": {1: "K"}, "under_light": {1: "K"}, "under_dark": {1: "K"}, "treat_logits": {1: "K"}, "role_logits": {1: "K"}},
                      opset_version=17, do_constant_folding=True, dynamo=False)
    onnx.checker.check_model(elem)
    esess = ort.InferenceSession(elem, so, providers=["CPUExecutionProvider"])
    eparity, agree, ok, n = 0.0, 0, 0, 0
    T = float(ck.get("temperature", 1.0))
    for f, soft in ef:
        with torch.no_grad():
            tt, rr = em(*[torch.from_numpy(f[k]) for k in ELEM_IN])
        o = esess.run(None, f)
        eparity = max(eparity, float(np.abs(tt.numpy() - o[0]).max()), float(np.abs(rr.numpy() - o[1]).max()))
        pred = o[0][0].argmax(-1)
        agree += int((pred == soft.argmax(-1)).sum()); ok += int((soft[np.arange(len(pred)), pred] > 0).sum()); n += len(pred)
    report = {"params": sum(p.numel() for p in model.parameters()),
              "elem_params": sum(p.numel() for k, p in model.named_parameters() if k.startswith(("cnn.", "img_in.", "elem_head."))),
              "parity_max_abs": parity, "elem_parity_max_abs": eparity, "temperature": T, "treat": TREAT, "img_roles": IMG_ROLES,
              "elem_val_sample": {"n": n, "acc": agree / max(1, n), "ok": ok / max(1, n)}, "variants": {}}
    print("parity", parity, "elem parity", eparity, flush=True)
    r = fidelity(ref, ref, val[:50]); r["size_mb"] = round(os.path.getsize(fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, [feeds_for(s) for s in val[:20]]), 2)
    report["variants"]["fp32"] = r
    report["elem"] = {"size_mb": round(os.path.getsize(elem) / 1e6, 2), "cpu_ms_median": round(bench(esess, [f for f, _ in ef]), 2) if ef else None}
    print(report, flush=True)
    json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1)


if __name__ == "__main__":
    main()