#!/usr/bin/env python3
import torch
from ultralytics import YOLO
from ultralytics.nn.modules.head import Detect
import os
import shutil


def npu_detect_forward(self, x):
    """
    YOLOv8 Detect Head Modified for NPU.

    YOLOv8 Detect structure (NOT end2end):
        - cv2: box regression layers (DFL, out channels = 4 * reg_max = 64)
        - cv3: classification layers (out channels = nc, nc=80 for COCO)

    Output:
        List of Tensors (6 items for 3 scales), Layout: NHWC.
        For each scale:
          - Box_Raw (B, H, W, 4*reg_max),  <-- raw DFL box logits (decode in post-process)
          - Cls_Raw (B, H, W, nc),         <-- class scores (raw logits), nc=80
    """
    if not isinstance(x, (list, tuple)):
        x = [x]

    res = []

    box_layers = self.cv2
    cls_layers = self.cv3

    for i in range(self.nl):
        # 1. Box branch (raw DFL logits) - NHWC
        bboxes = box_layers[i](x[i]).permute(0, 2, 3, 1)

        # 2. Cls branch (raw logits) - NHWC
        scores = cls_layers[i](x[i]).permute(0, 2, 3, 1)

        res.append(bboxes)
        res.append(scores)

    return res


def batch_export_yolov8_det():
    variants = ['n', 's', 'm', 'l', 'x']
    imgsz = 640

    # Execute Monkey Patch
    Detect.forward = npu_detect_forward
    print("Monkey patch applied for Detect: Output Layout forced to NHWC (Box, Cls for each scale).")

    for v in variants:
        model_name = f"yolov8{v}"
        pt_path = f"{model_name}.pt"
        onnx_final_name = f"{model_name}_640x640.onnx"
        print(f"\n--- Processing {model_name} ---")
        try:
            # Load model
            model = YOLO(pt_path)

            # Reapply monkey patch
            Detect.forward = npu_detect_forward

            # Ensure the model's head also uses the new forward
            if hasattr(model.model, 'model') and len(model.model.model) > 0:
                head = model.model.model[-1]
                if isinstance(head, Detect):
                    head.forward = lambda x: npu_detect_forward(head, x)

            # Execute export
            exported_path = model.export(
                format="onnx",
                imgsz=imgsz,
                dynamic=False,
                opset=11,
                simplify=True,
                nms=False
            )

            # Move and rename
            if exported_path:
                shutil.move(exported_path, onnx_final_name)
                print(f"Success: {onnx_final_name}")
        except Exception as e:
            print(f"Failed to export {model_name}: {e}")
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    batch_export_yolov8_det()