#!/usr/bin/env python3
"""Export task heads that have no Model Zoo example and compile them with RKNN."""

import argparse
import hashlib
import os
from pathlib import Path
import sys
import tempfile


TASKS = {
    "yolov5n-cls": "classify",
    "yolov8n-cls": "classify",
    "yolo11n-cls": "classify",
}
YOLO26_TASKS = {"cls": "classify", "seg": "segment", "pose": "pose", "obb": "obb",
                "sem": "semantic", "depth": "depth"}
TASKS.update({f"{family}{size}-cls": "classify"
              for family in ("yolov5", "yolov8", "yolo11") for size in "sm"})
TASKS.update({f"yolo26{size}-{suffix}": task
              for size in "nsm" for suffix, task in YOLO26_TASKS.items()})
SIZES = {"classify": 224, "semantic": 512, "depth": 384}
CHIPS = ("rk3576", "rk3588")
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def load_yolov5_classifier(checkpoint, source):
    """Bake YOLOv5's ImageNet normalization into the input graph."""
    import torch

    sys.path.insert(0, str(source))
    try:
        from models.experimental import attempt_load

        classifier = attempt_load(str(checkpoint), device=torch.device("cpu"), fuse=True)
    finally:
        sys.path.pop(0)

    class NormalizeClassifier(torch.nn.Module):
        def __init__(self, model):
            super().__init__()
            self.model = model
            self.register_buffer("mean", torch.tensor(IMAGENET_MEAN).reshape(1, 3, 1, 1))
            self.register_buffer("std", torch.tensor(IMAGENET_STD).reshape(1, 3, 1, 1))

        def forward(self, images):
            return self.model((images - self.mean) / self.std).softmax(dim=1)

    return NormalizeClassifier(classifier).eval(), classifier.names


def split_yolo26(head, task, size):
    """Expose one-to-one heads before TopK, following the Zoo YOLO26 detection layout."""
    import torch
    from ultralytics.utils.tal import make_anchors

    def forward(features):
        predictions = head.forward_head(features, **head.one2one)
        anchors, strides = make_anchors(features, head.stride, 0.5)
        head.anchors = anchors.transpose(0, 1)
        head.strides = strides.transpose(0, 1)
        raw_boxes = predictions["boxes"]
        left_top, right_bottom = raw_boxes.split(2, dim=1)
        anchor = head.anchors.unsqueeze(0)
        if task == "obb":
            angle = predictions["angle"]
            offset_x, offset_y = ((right_bottom - left_top) / 2).split(1, dim=1)
            cosine, sine = angle.cos(), angle.sin()
            rotated = torch.cat((offset_x * cosine - offset_y * sine,
                                 offset_x * sine + offset_y * cosine), dim=1)
            boxes = torch.cat((anchor + rotated, left_top + right_bottom), dim=1)
        else:
            boxes = torch.cat((anchor - left_top, anchor + right_bottom), dim=1)

        outputs = (predictions["scores"].sigmoid(), boxes * (head.strides / size))
        if task == "segment":
            return (*outputs, predictions["mask_coefficient"], head.proto(features))
        if task == "obb":
            return (*outputs, angle)

        points = head.kpts_decode(predictions["kpts"]).reshape(1, *head.kpt_shape, -1)
        # Keep visibility and x/y on comparable scales for INT8 quantization.
        points = torch.cat((points[:, :, :2] / size, points[:, :, 2:]), dim=2)
        return (*outputs, points.reshape(1, head.nk, -1))

    return forward


def export_checkpoint(checkpoint, destination, size, yolov5_source):
    os.environ.setdefault("YOLO_AUTOINSTALL", "false")
    import onnx
    import torch
    import ultralytics
    import yaml
    from ultralytics import YOLO

    task = TASKS[checkpoint.stem]
    if checkpoint.stem.startswith("yolov5"):
        model, names = load_yolov5_classifier(checkpoint, yolov5_source)
        head = None
    else:
        model = YOLO(str(checkpoint), task=task).model.cpu().float().eval()
        head = model.model[-1]
        names = getattr(model, "names", None) or {0: "depth"}
    split = task in {"segment", "pose", "obb"}
    if split:
        if head.reg_max != 1 or getattr(head, "one2one_cv2", None) is None:
            raise ValueError("Expected a DFL-free, one-to-one YOLO26 head")
        # Fusion removes whichever trained branch is inactive.
        head.end2end = True
    if head is not None:
        model.fuse(verbose=False)
    with checkpoint.open("rb") as file:
        digest = hashlib.file_digest(file, "sha256").hexdigest()
    metadata = {
        "task": task,
        "batch": 1,
        "channels": 3,
        "imgsz": [size, size],
        "names": names,
        "preprocess": "center_crop" if task == "classify" else "letterbox",
        "coordinates": "pixels",
        "end2end": False,
        "export_source": "model-zoo-inspired-pytorch",
        "checkpoint": checkpoint.name,
        "checkpoint_sha256": digest,
        "ultralytics": ultralytics.__version__,
        "pytorch": str(torch.__version__),
    }
    if split:
        head.export = True
        head.forward = split_yolo26(head, task, size)
        metadata.update(coordinates="normalized", end2end=True)
        output_names = ["Sigmoid_scores", "boxes"]
        if task == "segment":
            output_names.extend(("mask_coefficients", "prototypes"))
        elif task == "pose":
            output_names.append("keypoints")
            metadata["kpt_shape"] = list(head.kpt_shape)
        else:
            output_names.append("angle")
    elif task == "classify":
        if head is not None:
            head.export = True
        else:
            metadata["checkpoint_source"] = (
                f"https://github.com/ultralytics/yolov5/releases/download/v7.0/{checkpoint.name}"
            )
            metadata["normalization"] = "ImageNet mean/std baked into ONNX"
        output_names = ["probabilities"]
    else:
        # Evaluation mode emits the native grid before terminal full-image resize.
        head.export = False
        metadata["dense_resize"] = "half_pixel" if task == "semantic" else "align_corners"
        output_names = ["logits" if task == "semantic" else "depth"]

    image = torch.zeros(1, 3, size, size)
    with torch.inference_mode():
        expected = model(image)
        tensors = expected if isinstance(expected, tuple) else (expected,)
        if len(tensors) != len(output_names) or not all(torch.is_tensor(x) for x in tensors):
            raise ValueError("Unexpected output contract from checkpoint head")
        for name, tensor in zip(output_names, tensors):
            print(f"{name}: {tuple(tensor.shape)}")
        torch.onnx.export(model, image, str(destination / "model.onnx"), opset_version=13,
                          input_names=["images"], output_names=output_names,
                          dynamo=False, do_constant_folding=True)

    graph = onnx.load(destination / "model.onnx")
    onnx.checker.check_model(graph)
    (destination / "metadata.yaml").write_text(yaml.safe_dump(metadata, sort_keys=False))


def compile_graph(source, destination, chip, precision, dataset):
    import yaml
    from rknn.api import RKNN

    metadata = yaml.safe_load((source / "metadata.yaml").read_text())
    metadata.update(target_platform=chip, precision="int8" if precision == "i8" else "fp16")

    def check(code, step):
        if code != 0:
            raise RuntimeError(f"RKNN {step} failed ({code})")

    toolkit = RKNN(verbose=False)
    try:
        check(toolkit.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],
                             target_platform=chip), "config")
        check(toolkit.load_onnx(model=str(source / "model.onnx")), "load ONNX")
        with tempfile.TemporaryDirectory(prefix="nvr-calibration-") as directory:
            if precision == "i8" and metadata["checkpoint"].startswith("yolov5"):
                dataset = center_crop_calibration(dataset, Path(directory), metadata["imgsz"][0])
            check(toolkit.build(do_quantization=precision == "i8",
                                dataset=str(dataset) if dataset else None), "build")
        check(toolkit.export_rknn(str(destination / "model.rknn")), "export")
    finally:
        toolkit.release()
    (destination / "metadata.yaml").write_text(yaml.safe_dump(metadata, sort_keys=False))


def center_crop_calibration(dataset, directory, size):
    """Use the same square crop and linear resize as YOLOv5 classify/predict.py."""
    import cv2

    paths = []
    for index, line in enumerate(dataset.read_text().splitlines()):
        if not line.strip():
            continue
        source = dataset.parent / line.strip()
        image = cv2.imread(str(source), cv2.IMREAD_COLOR)
        if image is None:
            raise ValueError(f"Cannot read calibration image: {source}")
        height, width = image.shape[:2]
        side = min(height, width)
        top, left = (height - side) // 2, (width - side) // 2
        image = cv2.resize(image[top:top + side, left:left + side], (size, size))
        output = directory / f"{index:04}.png"
        if not cv2.imwrite(str(output), image):
            raise RuntimeError(f"Cannot write calibration image: {output}")
        paths.append(str(output))
    cropped = directory / "images.txt"
    cropped.write_text("\n".join(paths) + "\n")
    return cropped


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    export = commands.add_parser("export", help="Export one original checkpoint to ONNX")
    export.add_argument("checkpoint", type=Path)
    export.add_argument("output", type=Path, help="New directory for the ONNX and metadata")
    export.add_argument("--imgsz", type=int, help="Override the square input size")
    export.add_argument("--yolov5-source", type=Path, help="YOLOv5 v7.0 source checkout")
    compile_parser = commands.add_parser("compile", help="Build one RKNN from an exported directory")
    compile_parser.add_argument("onnx_directory", type=Path)
    compile_parser.add_argument("output", type=Path, help="New directory for the RKNN and metadata")
    compile_parser.add_argument("--chip", required=True, choices=CHIPS)
    compile_parser.add_argument("--precision", required=True, choices=("fp", "i8"))
    compile_parser.add_argument("--dataset", type=Path, help="RKNN calibration image list")
    args = parser.parse_args()

    destination = args.output.resolve()
    if destination.exists():
        parser.error(f"Output already exists: {destination}")
    destination.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix=".convert-", dir=destination.parent) as directory:
        work = Path(directory) / "output"
        work.mkdir()
        if args.command == "export":
            checkpoint = args.checkpoint.resolve()
            if not checkpoint.is_file() or checkpoint.stem not in TASKS:
                parser.error("Provide one of these local checkpoints: " + ", ".join(TASKS))
            size = args.imgsz or SIZES.get(TASKS[checkpoint.stem], 640)
            if size < 32 or size > 4096 or size % 32:
                parser.error("--imgsz must be a multiple of 32 from 32 to 4096")
            yolov5_source = args.yolov5_source.resolve() if args.yolov5_source else None
            if checkpoint.stem.startswith("yolov5") and (
                yolov5_source is None or not (yolov5_source / "models/experimental.py").is_file()
            ):
                parser.error("YOLOv5 classification requires a --yolov5-source v7.0 checkout")
            export_checkpoint(checkpoint, work, size, yolov5_source)
        else:
            dataset = args.dataset.resolve() if args.dataset else None
            if args.precision == "i8" and (dataset is None or not dataset.is_file()):
                parser.error("INT8 requires an existing --dataset image list")
            compile_graph(args.onnx_directory.resolve(), work, args.chip, args.precision, dataset)
        work.rename(destination)
    print(destination)


if __name__ == "__main__":
    main()