"""Export the RF-DETR Medium licence-plate checkpoint to ONNX at 576x576.

The checkpoint is our own fine-tune, made with `rfdetr==1.10.1` from
Roboflow's rf-detr-medium.pth on Open Images V7's "Vehicle registration plate"
boxes (annotations CC BY 4.0), the group-of boxes dropped: 4,497 training
images with 6,554 plates, 724 validation images with 985. Twenty epochs at
576, batch 8 with two steps of gradient accumulation, on one RTX 4090; the
EMA weights with the best validation mAP, epoch 4 at 0.655 mAP 50:95, are
the checkpoint exported here.

The export mirrors scripts/export_face.py: the rfdetr package's own
model.export(format="onnx"), opset 17. num_classes is NOT passed, so the head
is exported as trained. rfdetr 1.10.1 lays a one-class head out as the class
plus one spare slot, so `labels` is 1x300x2: index 0 is the plate, index 1
was never trained toward anything and stays quiet.
"""

import glob
import hashlib
import os
import shutil
import sys

from rfdetr import RFDETRMedium

CKPT = r"E:/projects/model-work/ckpt/rfdetr_medium_plate.pth"
OUTDIR = r"E:/projects/model-work/export-plate"
FINAL = r"E:/projects/rfdetr-onnx/rf-detr-medium-plate.onnx"


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main():
    os.makedirs(OUTDIR, exist_ok=True)
    model = RFDETRMedium(pretrain_weights=CKPT, resolution=576)
    print("model_config.num_classes =", model.model_config.num_classes)
    print("model_config.resolution  =", model.model_config.resolution)
    print("class_embed out_features =", model.model.model.class_embed.out_features)
    model.export(format="onnx", output_dir=OUTDIR, opset_version=17)
    produced = sorted(glob.glob(os.path.join(OUTDIR, "*.onnx")))
    for path in produced:
        print("produced", path, os.path.getsize(path), sha256(path))
    if len(produced) != 1:
        print("expected exactly one .onnx, got", produced)
        sys.exit(1)
    if os.path.exists(FINAL):
        print("refusing to overwrite existing", FINAL)
        sys.exit(1)
    shutil.copyfile(produced[0], FINAL)
    print("final", FINAL, os.path.getsize(FINAL), sha256(FINAL))
    print("source_pth", CKPT, os.path.getsize(CKPT), sha256(CKPT))


if __name__ == "__main__":
    main()