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


def npu_pose_forward(self, x):
    """
    YOLOv8 Pose Head Modified for NPU.

    YOLOv8 Pose structure (NOT end2end):
        - cv2: box regression layers (DFL, out channels = 4 * reg_max = 64)
        - cv3: classification layers (out channels = nc, nc=1 for pose)
        - cv4: keypoint output layers (out channels = nk = 17 * 3 = 51)

    Output:
        List of Tensors (9 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=1 for pose
          - Kpt_Raw (B, H, W, nk),         <-- raw keypoints (nk = 17*3 = 51)
    """
    if not isinstance(x, (list, tuple)):
        x = [x]

    res = []

    box_layers = self.cv2
    cls_layers = self.cv3
    kpt_layers = self.cv4

    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)

        # 3. Keypoints branch (raw) - NHWC
        kpts = kpt_layers[i](x[i]).permute(0, 2, 3, 1)

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

    return res


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

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

    for v in variants:
        model_name = f"yolov8{v}-pose"
        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
            Pose.forward = npu_pose_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, Pose):
                    head.forward = lambda x: npu_pose_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_pose()