#!/usr/bin/env python3
"""Convert the exported YOLOv5 ONNX model to an INT8 RK3588 RKNN model.

Run this script on an x86_64 Ubuntu 20.04 / Python 3.8 machine with
RKNN Toolkit2 1.5.2 installed.  It is deliberately not run on the robot;
the robot only needs RKNN Toolkit Lite2 to execute the resulting file.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from rknn.api import RKNN


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--onnx", required=True, type=Path, help="Path to best.onnx")
    parser.add_argument(
        "--dataset",
        required=True,
        type=Path,
        help="Text file containing absolute paths to representative PNG/JPG images",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("traffic_sign_direction.rknn"),
        help="Output RKNN model path",
    )
    parser.add_argument(
        "--no-quantize",
        action="store_true",
        help="Build an FP model for conversion diagnosis (larger and slower than INT8)",
    )
    return parser.parse_args()


def check_inputs(onnx_path: Path, dataset_path: Path, needs_calibration: bool) -> None:
    if not onnx_path.is_file():
        raise FileNotFoundError("ONNX model not found: %s" % onnx_path)
    if not needs_calibration:
        return
    if not dataset_path.is_file():
        raise FileNotFoundError("Calibration dataset list not found: %s" % dataset_path)

    images = [Path(line.strip()) for line in dataset_path.read_text(encoding="utf-8").splitlines() if line.strip()]
    if len(images) < 20:
        raise ValueError("Use at least 20 representative images for INT8 calibration; got %d" % len(images))
    missing = [str(image) for image in images if not image.is_file()]
    if missing:
        raise FileNotFoundError("Missing image paths in %s, first: %s" % (dataset_path, missing[0]))
    print("INT8 calibration images: %d" % len(images))


def main() -> int:
    args = parse_args()
    onnx_path = args.onnx.expanduser().resolve()
    dataset_path = args.dataset.expanduser().resolve()
    output_path = args.output.expanduser().resolve()
    quantize = not args.no_quantize
    check_inputs(onnx_path, dataset_path, needs_calibration=quantize)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    rknn = RKNN(verbose=True)
    try:
        # The ROS node supplies a 640x640 RGB uint8 image.  RKNN performs the
        # equivalent of YOLOv5's /255 input normalization inside the graph.
        result = rknn.config(
            target_platform="rk3588",
            mean_values=[[0, 0, 0]],
            std_values=[[255, 255, 255]],
            optimization_level=3,
        )
        if result != 0:
            raise RuntimeError("rknn.config failed: %s" % result)
        result = rknn.load_onnx(model=str(onnx_path))
        if result != 0:
            raise RuntimeError("rknn.load_onnx failed: %s" % result)
        build_options = {"do_quantization": quantize}
        if quantize:
            build_options["dataset"] = str(dataset_path)
        result = rknn.build(**build_options)
        if result != 0:
            raise RuntimeError("rknn.build failed: %s" % result)
        result = rknn.export_rknn(str(output_path))
        if result != 0:
            raise RuntimeError("rknn.export_rknn failed: %s" % result)
    finally:
        rknn.release()

    print("RKNN conversion complete (%s): %s" % ("INT8" if quantize else "FP", output_path))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (FileNotFoundError, RuntimeError, ValueError) as error:
        print("ERROR: %s" % error, file=sys.stderr)
        raise SystemExit(1)