"""ONNX (RT-DETRv4-S / D-FINE, ogkalu/comic-text-and-bubble-detector) -> Core ML.

Steps
 1. Cut the exported ONNX before its post-processor (TopK / label decode / orig-size scaling) so the
    Core ML model returns raw per-query outputs:
       logits [1,300,3]  (pre-sigmoid class logits: bubble, text_bubble, text_free)
       boxes  [1,300,4]  (cx, cy, w, h normalised to the 640x640 input)
 2. Patch Reshape(allowzero=1) -> 0 (no zero-sized dims in this graph) and register the few
    onnx2torch converters missing for opset 18 (Resize-18, Split-18, ReduceMean/Max-18, GridSample-16).
 3. onnx2torch -> torch.jit.trace (fixed 1x3x640x640) -> coremltools mlprogram, fp16 and fp32
    (+ optional int8 weight-only palettisation/linear quantisation).
 4. Parity vs onnxruntime on the benchmark pages.

Usage: python convert_v4s_coreml.py <src.onnx> <out_prefix>
   e.g. python convert_v4s_coreml.py models/ogkalu-ctbd/detector-v4-s.fp32-5084e7d.onnx coreml/detector-chosen/ogkalu-v4s
"""
import sys, os, glob
import numpy as np, onnx, torch, torch.nn.functional as F
from onnx.utils import extract_model
from onnx2torch import convert
from onnx2torch.node_converters.registry import add_converter, _CONVERTER_REGISTRY, OperationDescription
from onnx2torch.utils.common import OperationConverterResult, onnx_mapping_from_node, get_const_value, OnnxToTorchModule

import os
B = os.environ.get("NEMU_OCR_WORKDIR", ".")  # holds models/ (upstream weights), coreml/ (outputs)
SRC, PREFIX = sys.argv[1], sys.argv[2]
LOGITS, BOXES = sys.argv[3] if len(sys.argv) > 3 else "select_36", sys.argv[4] if len(sys.argv) > 4 else "select_37"


def reg(op, ver, fn):
    k = OperationDescription(domain="", operation_type=op, version=ver)
    if k not in _CONVERTER_REGISTRY: add_converter(operation_type=op, version=ver)(fn)


get = lambda op, v: _CONVERTER_REGISTRY[OperationDescription(domain="", operation_type=op, version=v)]
reg("Resize", 18, get("Resize", 13))
reg("Resize", 19, get("Resize", 13))
reg("Split", 18, get("Split", 13))


class ReduceStatic(torch.nn.Module, OnnxToTorchModule):
    def __init__(self, fn, axes, keepdims):
        super().__init__(); self.fn, self.axes, self.keepdims = fn, axes, bool(keepdims)

    def forward(self, x, *_):
        if self.fn == "mean": return x.mean(dim=self.axes, keepdim=self.keepdims)
        out = x
        for a in sorted([a % x.dim() for a in self.axes], reverse=True): out = out.amax(dim=a, keepdim=self.keepdims)
        return out


def reduce_conv(fn):
    def _(node, graph):
        axes = node.attributes.get("axes")
        if axes is None and len(node.input_values) > 1:
            axes = get_const_value(node.input_values[1], graph).tolist()
        return OperationConverterResult(torch_module=ReduceStatic(fn, list(axes), node.attributes.get("keepdims", 1)),
                                        onnx_mapping=onnx_mapping_from_node(node))
    return _


reg("ReduceMean", 18, reduce_conv("mean"))
reg("ReduceMax", 18, reduce_conv("max"))


class GridSample(torch.nn.Module, OnnxToTorchModule):
    def __init__(self, mode, padding_mode, align_corners):
        super().__init__(); self.mode, self.padding_mode, self.align_corners = mode, padding_mode, align_corners

    def forward(self, x, grid):
        return F.grid_sample(x, grid, mode=self.mode, padding_mode=self.padding_mode, align_corners=self.align_corners)


def gs(node, graph):
    a = node.attributes
    mode = {"linear": "bilinear", "bilinear": "bilinear", "nearest": "nearest", "cubic": "bicubic", "bicubic": "bicubic"}[a.get("mode", "linear")]
    return OperationConverterResult(torch_module=GridSample(mode, a.get("padding_mode", "zeros"), bool(a.get("align_corners", 0))),
                                    onnx_mapping=onnx_mapping_from_node(node))


reg("GridSample", 16, gs)
reg("GridSample", 20, gs)

# 1+2: cut + patch
core = f"{B}/detector/convert/{os.path.basename(PREFIX)}-core.onnx"
extract_model(SRC, core, ["images"], [LOGITS, BOXES])
m = onnx.load(core)
for n in m.graph.node:
    if n.op_type == "Reshape":
        for at in n.attribute:
            if at.name == "allowzero": at.i = 0
onnx.save(m, core)
tm = convert(m).eval()

# reference: onnxruntime on the cut graph
import onnxruntime as ort
from PIL import Image
sess = ort.InferenceSession(core, providers=["CPUExecutionProvider"])
def prep(f):
    return (np.asarray(Image.open(f).convert("RGB").resize((640, 640), Image.BILINEAR), dtype=np.float32).transpose(2, 0, 1)[None] / 255.0)
x = prep(f"{B}/pages/hanako-ch1-p02.png")
ref = sess.run(None, {"images": x})
with torch.no_grad(): t = tm(torch.from_numpy(x))
print("onnx2torch vs ORT max abs diff:", [float(np.abs(r - tt.numpy()).max()) for r, tt in zip(ref, t)])


class Wrap(torch.nn.Module):
    def __init__(self, m): super().__init__(); self.m = m
    def forward(self, images):
        logits, boxes = self.m(images)
        return logits, boxes


with torch.no_grad():
    traced = torch.jit.trace(Wrap(tm).eval(), torch.from_numpy(x), check_trace=False)

import coremltools as ct
for prec in ("fp32", "fp16"):
    ml = ct.convert(traced, convert_to="mlprogram",
                    inputs=[ct.TensorType(name="images", shape=(1, 3, 640, 640), dtype=np.float32)],
                    outputs=[ct.TensorType(name="logits"), ct.TensorType(name="boxes")],
                    compute_precision=ct.precision.FLOAT16 if prec == "fp16" else ct.precision.FLOAT32,
                    minimum_deployment_target=ct.target.iOS17)
    ml.short_description = ("ogkalu/comic-text-and-bubble-detector detector-v4-s (RT-DETRv4-S / D-FINE decoder). "
                            "Input: RGB page resized (no letterbox) to 640x640, NCHW float32 in [0,1]. "
                            "Outputs: 300 queries x 3 class logits (bubble, text_bubble, text_free; apply sigmoid) "
                            "and cx,cy,w,h boxes normalised to the page. See postprocess-spec.md.")
    ml.author = "ogkalu (weights, Apache-2.0); Core ML conversion for nemu"
    ml.license = "Apache-2.0 (https://huggingface.co/ogkalu/comic-text-and-bubble-detector)"
    ml.version = "v4-s int8 @ 16e8a622 dequantised" if "deq" in SRC else "v4-s fp32 @ 5084e7d"
    ml.user_defined_metadata["classes"] = "bubble,text_bubble,text_free"
    ml.user_defined_metadata["input_size"] = "640x640"
    ml.user_defined_metadata["source_onnx"] = os.path.basename(SRC)
    ml.input_description["images"] = "1x3x640x640 float32 RGB in [0,1]; page resized with antialiased bilinear (PIL BILINEAR), no letterbox, no mean/std"
    ml.output_description["logits"] = "1x300x3 pre-sigmoid logits: bubble, text_bubble, text_free"
    ml.output_description["boxes"] = "1x300x4 (cx, cy, w, h) normalised to [0,1] of the original page"
    ml.save(f"{PREFIX}-{prec}.mlpackage")
    out = ml.predict({"images": x})
    print(prec, "Core ML vs ORT max abs diff: logits", float(np.abs(out["logits"] - ref[0]).max()), "boxes", float(np.abs(out["boxes"] - ref[1]).max()))