import torch
from ultralytics import YOLO
from ultralytics.nn.modules.head import OBB26
import os
import shutil


def npu_obb26_forward(self, x):
    """
    YOLO26 OBB26 Head Modified for NPU.

    OBB26 structure:
        - cv2 / one2one_cv2: box regression layers (output 4*reg_max channels, reg_max=1 -> 4)
        - cv3 / one2one_cv3: classification layers (output nc channels, nc=15 for DOTAv1)
        - cv4 / one2one_cv4: angle regression layers (output ne channels, ne=1)

    Note:
        OBB26 outputs raw angle predictions (in radians) without sigmoid transformation,
        which is different from the original OBB head where the angle is wrapped via
        ``(angle.sigmoid() - 0.25) * pi`` in ``forward_head``. The post-process therefore
        consumes the raw angle directly via ``cos(angle)`` / ``sin(angle)``.

    Output:
        List of Tensors (9 items: 3 scales * 3 outputs), Layout: NHWC for all branches.
        For each scale:
          - Box_Raw   (B, H, W, 4*reg_max),  <-- bbox regression (ltrb distance, reg_max=1)
          - Cls_Raw   (B, H, W, nc),         <-- class scores (raw logits)
          - Angle_Raw (B, H, W, ne),         <-- rotation angle (raw radians, ne=1)
    """
    if not isinstance(x, (list, tuple)):
        x = [x]

    res = []

    # Box / Cls use one2one branch (end2end mode) when available
    if hasattr(self, 'one2one_cv2') and hasattr(self, 'one2one_cv3'):
        box_layers = self.one2one_cv2
        cls_layers = self.one2one_cv3
    else:
        box_layers = self.cv2
        cls_layers = self.cv3

    # Angle uses one2one_cv4 in end2end mode
    if hasattr(self, 'one2one_cv4'):
        angle_layers = self.one2one_cv4
    else:
        angle_layers = self.cv4

    for i in range(self.nl):
        # 1. Box branch - NHWC (4 * reg_max channels, reg_max=1 -> 4)
        bboxes = box_layers[i](x[i]).permute(0, 2, 3, 1)

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

        # 3. Angle branch - NHWC (ne channels, raw radians)
        angle = angle_layers[i](x[i]).permute(0, 2, 3, 1)

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

    return res


def batch_export_yolo26_obb():
    variants = ['n', 's', 'm', 'l', 'x']
    imgsz = 1024  # YOLO26-OBB officially uses 1024x1024 (DOTAv1 pretrained)

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

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

            # Reapply monkey patch
            OBB26.forward = npu_obb26_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, OBB26):
                    head.forward = lambda x: npu_obb26_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_yolo26_obb()