# 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|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 recognizer's Reshapes compute their target shape from tensor sizes, which tracing freezes at the
# traced width; 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.

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 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()
    det = convert(onnx.load(onnx_path), 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(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()