# Converts a PP-OCRv5 detector and recognizer to Core ML models for modules/OCR on macOS, then checks
# each against the ONNX model on the CPU. Run once, on a Mac, after modules/OCR/fetch_models.jai;
# mcap only uses the .mlpackage files it writes (into modules/OCR/models, beside the ONNX ones).
#
#   python3 -m venv ~/mcap-coreml
#   ~/mcap-coreml/bin/pip install torch==2.7.0 coremltools==9.0 onnx onnxruntime pillow
#   ~/mcap-coreml/bin/pip install --no-deps onnx2torch torchvision==0.22.0
#   ~/mcap-coreml/bin/python tools/export_ppocr_coreml.py --model mobile
#   ~/mcap-coreml/bin/python tools/export_ppocr_coreml.py --model server
#
# There is no PyTorch source for PP-OCR (it is PaddlePaddle's), and coremltools no longer reads ONNX, so
# each model goes ONNX -> PyTorch (onnx2torch) -> Core ML, as an ML Program in float32: float16 moves
# the detector's text pixels and changes recognized text, and the Neural Engine gets it wrong.
#
# The shapes. Core ML on the GPU cannot take a flexible height and width for the detector (its
# depthwise convolutions fail to compile: "invalid axis" in mps.expand_dims), so both models take a
# list of sizes (at most 128 a model), which mcap pads its inputs up to; models.config's coreml_ keys
# say the same steps, and must be kept in step with these:
#   detector:   1 x 3 x H x W, H a multiple of 128 up to 1536 (det_tile_height), W of 256 up to 2560 (det_max_side)
#   recognizer: 16 x 3 x 48 x W, W a multiple of 64 up to 4096 (REC_MAX_WIDTH); fewer lines are padded to 16
# The server models pad some 2 x 2 convolutions and max pools with auto_pad SAME_UPPER, which onnx2torch
# does not take; at stride 1 it is a fixed pixel after, written out as pads first.
# Reshapes that compute their target shape from tensor sizes (the recognizers', the server detector's)
# are frozen by tracing at the traced size; they are rewritten first as constant shapes of 0 (copy the
# dimension) and -1 (infer it), worked out from the shapes onnxruntime gives at several sizes, and
# checked to give the same output.
#
# Measured on an M5 Pro against onnxruntime on the CPU, each at the same (padded) size: the detector
# about 5 times as fast, the recognizer about 4 times for 16 lines, with the same text pixels and text.
# Whole images through mcap (ocr-bench, 8 text pages): mobile 323 ms an image against 536, server 940
# against 2030.

import argparse, itertools, os, sys, time
import numpy as np
import onnx, onnxruntime as ort
from onnx import helper, numpy_helper
import torch, onnx2torch
import coremltools as ct

HERE = os.path.dirname(os.path.abspath(__file__))
MODELS = os.path.join(HERE, "..", "modules", "OCR", "models")

DET_HEIGHTS = list(range(128, 1536 + 1, 128))
DET_WIDTHS  = list(range(256, 2560 + 1, 256))
REC_BATCH   = 16
REC_HEIGHT  = 48
REC_WIDTHS  = list(range(64, 4096 + 1, 64))

# coremltools 9.0 casts a folded one-element tensor with int(array), which NumPy 2 refuses ("only
# 0-dimensional arrays can be converted"); the same cast, with the element taken out first.
from coremltools.converters.mil.frontend.torch import ops as ct_ops
_original_cast = ct_ops._cast
def _cast_one_element(context, node, dtype, dtype_name):
    x = ct_ops._get_inputs(context, node, expected=1)[0]
    if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
        context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
        return
    _original_cast(context, node, dtype, dtype_name)
ct_ops._cast = _cast_one_element


def constant_reshapes(model, input_shapes):
    """The model with each computed Reshape target made a constant of 0s and -1s."""
    graph = model.graph
    initializers = {i.name for i in graph.initializer}
    producers = {o: n for n in graph.node for o in n.output}
    computed = [n for n in graph.node if n.op_type == "Reshape" and n.input[1] not in initializers
                and not (n.input[1] in producers and producers[n.input[1]].op_type == "Constant")]
    if not computed:
        return model
    probe = onnx.ModelProto(); probe.CopyFrom(model)
    names = []
    for n in computed:
        for t in (n.input[0], n.output[0]):
            if t not in names:
                names.append(t)
                probe.graph.output.append(helper.make_tensor_value_info(t, onnx.TensorProto.FLOAT, None))
    session = ort.InferenceSession(probe.SerializeToString(), providers=["CPUExecutionProvider"])
    name = session.get_inputs()[0].name
    shapes = []
    for shape in input_shapes:
        outs = session.run(names, {name: np.zeros(shape, np.float32)})
        shapes.append({t: o.shape for t, o in zip(names, outs)})
    for n in computed:
        ins = [s[n.input[0]] for s in shapes]
        outs = [s[n.output[0]] for s in shapes]
        target = []
        for d in range(len(outs[0])):
            sizes = {o[d] for o in outs}
            if len(sizes) == 1:
                target.append(sizes.pop())
            elif all(len(i) > d and i[d] == o[d] for i, o in zip(ins, outs)):
                target.append(0)
            else:
                target.append(-1)
        if target.count(-1) > 1:
            sys.exit(f"{n.name}: more than one dimension to infer ({target})")
        constant = n.name + "_constant_shape"
        graph.initializer.append(numpy_helper.from_array(np.array(target, dtype=np.int64), constant))
        n.input[1] = constant
        print(f"  {n.name}: target shape {target}")
    return model


def explicit_padding(model):
    """The model with auto_pad SAME_UPPER and SAME_LOWER, which onnx2torch does not take, written out as
    pads. Only stride 1 is rewritten: there the padding does not depend on the input's size (the server
    models' 2 x 2 convolutions and max pools: one pixel after, none before)."""
    for n in model.graph.node:
        attrs = {a.name: a for a in n.attribute}
        mode = attrs["auto_pad"].s.decode() if "auto_pad" in attrs else "NOTSET"
        if mode not in ("SAME_UPPER", "SAME_LOWER"):
            continue
        kernel = list(attrs["kernel_shape"].ints) if "kernel_shape" in attrs else None
        if kernel is None:
            sys.exit(f"{n.name}: auto_pad {mode} with no kernel_shape")
        strides = list(attrs["strides"].ints) if "strides" in attrs else [1] * len(kernel)
        dilations = list(attrs["dilations"].ints) if "dilations" in attrs else [1] * len(kernel)
        if any(s != 1 for s in strides):
            sys.exit(f"{n.name}: auto_pad {mode} with strides {strides}; only stride 1 is rewritten")
        totals = [(k - 1) * d for k, d in zip(kernel, dilations)]
        small = [t // 2 for t in totals]
        begins = small if mode == "SAME_UPPER" else [t - b for t, b in zip(totals, small)]
        ends = [t - b for t, b in zip(totals, begins)]
        n.attribute.remove(attrs["auto_pad"])
        if "pads" in attrs:
            n.attribute.remove(attrs["pads"])
        n.attribute.append(helper.make_attribute("pads", begins + ends))
        print(f"  {n.name} ({n.op_type}): auto_pad {mode} -> pads {begins + ends}")
    return model


def convert(onnx_model, example, shape, path):
    module = onnx2torch.convert(onnx_model).eval()
    with torch.no_grad():
        traced = torch.jit.trace(module, torch.from_numpy(example))
    ml = ct.convert(traced, inputs=[ct.TensorType(name="x", shape=shape)], outputs=[ct.TensorType(name="y")],
                    convert_to="mlprogram", compute_precision=ct.precision.FLOAT32,
                    minimum_deployment_target=ct.target.macOS14)
    ml.save(path)
    return ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_GPU)


def ctc(logits):
    texts = []
    for row in logits.argmax(-1):
        seq, last = [], -1
        for i in row:
            if i != last and i != 0:
                seq.append(int(i))
            last = i
        texts.append(seq)
    return texts


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", choices=["mobile", "server"], default="mobile")
    args = parser.parse_args()
    rng = np.random.default_rng(0)

    # Detector.
    stem = f"ch_PP-OCRv5_det_{args.model}"
    onnx_path = os.path.join(MODELS, stem + ".onnx")
    out_path = os.path.join(MODELS, stem + ".mlpackage")
    shapes = [(1, 3, h, w) for h in DET_HEIGHTS for w in DET_WIDTHS]
    start = time.perf_counter()
    model = constant_reshapes(explicit_padding(onnx.load(onnx_path)), [(1, 3, h, w) for h in (256, 512) for w in (512, 768)])
    det = convert(model, np.zeros((1, 3, 768, 1280), np.float32),
                  ct.EnumeratedShapes(shapes=shapes, default=(1, 3, 768, 1280)), out_path)
    print(f"Wrote {out_path} ({len(shapes)} sizes) in {time.perf_counter() - start:.1f} s.")
    reference = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    worst = 0.0
    for (h, w) in [(768, 1280), (1536, 2560), (128, 256)]:
        x = rng.standard_normal((1, 3, h, w)).astype(np.float32)
        a = reference.run(None, {"x": x})[0]
        b = det.predict({"x": x})["y"]
        flips = float(np.mean((a > 0.3) != (b > 0.3)))
        worst = max(worst, flips)
        print(f"  {h} x {w}: max difference {np.abs(a - b).max():.2e}, pixels across 0.3: {flips * 100:.4f}%")
    if worst > 0.0001:
        sys.exit("The Core ML detector is too far from the ONNX one.")

    # Recognizer.
    stem = f"ch_PP-OCRv5_rec_{args.model}"
    onnx_path = os.path.join(MODELS, stem + ".onnx")
    out_path = os.path.join(MODELS, stem + ".mlpackage")
    print("Recognizer: making its reshapes constant")
    model = constant_reshapes(explicit_padding(onnx.load(onnx_path)), [(n, 3, REC_HEIGHT, w) for n in (1, 2) for w in (320, 640)])
    shapes = [(REC_BATCH, 3, REC_HEIGHT, w) for w in REC_WIDTHS]
    start = time.perf_counter()
    rec = convert(model, np.zeros((REC_BATCH, 3, REC_HEIGHT, 640), np.float32),
                  ct.EnumeratedShapes(shapes=shapes, default=(REC_BATCH, 3, REC_HEIGHT, 640)), out_path)
    print(f"Wrote {out_path} ({len(shapes)} widths) in {time.perf_counter() - start:.1f} s.")
    reference = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    for w in (64, 640, 4096):
        x = rng.standard_normal((REC_BATCH, 3, REC_HEIGHT, w)).astype(np.float32) * 0.5
        a = reference.run(None, {"x": x})[0]
        b = rec.predict({"x": x})["y"]
        same = np.mean([p == q for p, q in zip(ctc(a), ctc(b))])
        print(f"  width {w}: max difference {np.abs(a - b).max():.2e}, decoded the same {same * 100:.0f}%")
        if same < 1:
            sys.exit("The Core ML recognizer reads differently from the ONNX one.")


if __name__ == "__main__":
    main()