"""Export trained components to ONNX for CPU inference with ONNX Runtime, and check parity.

    python -m tinydoc.export --weights runs --out models
"""

import argparse
from pathlib import Path

import numpy as np
import torch

from .features import MAX_CHARS, NUM_FEATS
from .models import build, count_params

OPSET = 17


def sample_inputs(name):
    g = torch.Generator().manual_seed(0)
    if name == "detector":
        return (torch.rand(1, 3, 640, 480, generator=g),), ["image"], ["kernel_logits"], {"image": {0: "batch", 2: "height", 3: "width"}}
    if name == "recognizer":
        x = torch.rand(2, 1, 32, 320, generator=g)
        mask = torch.zeros(2, 80, dtype=torch.bool)
        mask[1, 60:] = True
        return (x, mask), ["image", "step_mask"], ["logits"], {"image": {0: "batch", 3: "width"}, "step_mask": {0: "batch", 1: "steps"}}
    if name == "orientation":
        return (torch.rand(4, 1, 192, 192, generator=g),), ["patches"], ["logits"], {"patches": {0: "batch"}}
    n = 37
    return ((torch.randint(1, 100, (1, n, MAX_CHARS), generator=g), torch.rand(1, n, NUM_FEATS, generator=g),
             torch.rand(1, n, 4, generator=g) * 30, torch.ones(1, n)),
            ["char_ids", "feats", "rel_boxes", "valid"], ["field_logits", "rel_logits", "doc_logits"],
            {"char_ids": {0: "batch", 1: "segments"}, "feats": {0: "batch", 1: "segments"},
             "rel_boxes": {0: "batch", 1: "segments"}, "valid": {0: "batch", 1: "segments"}})


def alt_inputs(name):
    """Different shapes from the export sample, to prove the dynamic axes really are dynamic."""
    g = torch.Generator().manual_seed(1)
    if name == "detector":
        return (torch.rand(1, 3, 320, 704, generator=g),)
    if name == "recognizer":
        return torch.rand(3, 1, 32, 512, generator=g), torch.zeros(3, 128, dtype=torch.bool)
    if name == "orientation":
        return (torch.rand(7, 1, 192, 192, generator=g),)
    n = 120
    return (torch.randint(1, 100, (1, n, MAX_CHARS), generator=g), torch.rand(1, n, NUM_FEATS, generator=g),
            torch.rand(1, n, 4, generator=g) * 30, torch.ones(1, n))


def export_one(name, weights, out_dir, random_init=False):
    import onnxruntime as ort

    ck = None if random_init else torch.load(Path(weights) / name / "best.pt", map_location="cpu", weights_only=False)
    model = build(name, **((ck or {}).get("arch") or {}))
    if ck:
        model.load_state_dict(ck["model"])
    model.eval()
    args, in_names, out_names, dyn = sample_inputs(name)
    path = Path(out_dir) / f"{name}.onnx"
    torch.onnx.export(model, args, str(path), input_names=in_names, output_names=out_names,
                      dynamic_axes=dyn, opset_version=OPSET, dynamo=False)
    import onnx
    proto = onnx.load(str(path))
    meta = proto.metadata_props.add()
    meta.key, meta.value = "trained_parameters", str(count_params(model))
    onnx.save(proto, str(path))
    sess = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    worst = 0.0
    for inputs in (args, alt_inputs(name)):
        with torch.no_grad():
            ref = model(*inputs)
        ref = ref if isinstance(ref, tuple) else (ref,)
        got = sess.run(None, {n: t.numpy() for n, t in zip(in_names, inputs)})
        worst = max(worst, max(float(np.abs(r.numpy() - g).max()) for r, g in zip(ref, got)))
    size = path.stat().st_size / 1e6
    print(f"  {name:12s} params={count_params(model):>9,d}  {size:5.1f} MB  max |torch - onnx| = {worst:.2e}")
    return worst


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default="runs")
    ap.add_argument("--out", default="models")
    ap.add_argument("--random-init", action="store_true", help="export untrained weights (pipeline testing)")
    args = ap.parse_args()
    Path(args.out).mkdir(parents=True, exist_ok=True)
    worst = max(export_one(n, args.weights, args.out, args.random_init) for n in ("detector", "recognizer", "orientation", "extractor"))
    if worst > 1e-3:
        raise SystemExit(f"ONNX parity check failed (max diff {worst:.2e})")
    print(f"exported to {args.out}/ (parity ok)")


if __name__ == "__main__":
    main()