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


def npu_segment_forward(self, x):
    """
    YOLOv8 Segment Head Modified for NPU.

    YOLOv8 Segment structure (NOT end2end):
        - cv2: box regression layers (DFL, out channels = 4 * reg_max = 64)
        - cv3: classification layers (out channels = nc = 80)
        - cv4: mask coefficient layers (out channels = nm = 32)
        - proto: Proto module for generating mask prototypes

    Output:
        List of Tensors (10 items: 3 scales * 3 outputs + 1 proto), NHWC for detection, NCHW for proto
        [
          Scale1_Box_Raw (B, 80, 80, 4*reg_max),  <-- raw DFL box logits
          Scale1_Cls_Raw (B, 80, 80, nc),         <-- class scores (raw logits), nc=80
          Scale1_Mask_Raw (B, 80, 80, nm),        <-- mask coefficients, nm=32
          Scale2_Box_Raw (B, 40, 40, 4*reg_max),
          Scale2_Cls_Raw (B, 40, 40, nc),
          Scale2_Mask_Raw (B, 40, 40, nm),
          Scale3_Box_Raw (B, 20, 20, 4*reg_max),
          Scale3_Cls_Raw (B, 20, 20, nc),
          Scale3_Mask_Raw (B, 20, 20, nm),
          Proto (B, nm, H, W),                     <-- mask prototypes (160x160 for 640 input)
        ]
    """
    if not isinstance(x, (list, tuple)):
        x = [x]

    res = []

    box_layers = self.cv2
    cls_layers = self.cv3
    mask_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. Mask coefficients branch - NHWC
        masks = mask_layers[i](x[i]).permute(0, 2, 3, 1)

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

    # 4. Proto output - NCHW (keep original format for mask processing)
    proto = self.proto(x[0])
    res.append(proto)

    return res


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

    # Execute Monkey Patch
    Segment.forward = npu_segment_forward
    print("Monkey patch applied for Segment: Output Layout forced to NHWC + Proto.")

    for v in variants:
        model_name = f"yolov8{v}-seg"
        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
            Segment.forward = npu_segment_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, Segment):
                    head.forward = lambda x: npu_segment_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_seg()