#!/usr/bin/env python3
"""
Mac YOLO Inference Tool for armor detection
Independent of ROS2, uses ONNX Runtime with CoreML acceleration
"""

import argparse
import os
import sys
import time
from dataclasses import dataclass
from enum import Enum
from typing import List, Tuple, Optional

import cv2
import numpy as np
import onnxruntime as ort


class NumberClass(Enum):
    GUARD = 0       # G (哨兵)
    ONE = 1         # 1 (一号)
    TWO = 2         # 2 (二号)
    THREE = 3       # 3 (三号)
    FOUR = 4        # 4 (四号)
    FIVE = 5        # 5 (五号)
    OUTPOST = 6     # O (前哨站)
    BASE_SMALL = 7  # Bs (基地小装甲)
    BASE_LARGE = 8  # Bb (基地大装甲)
    NEGATIVE = 9    # negative


NUMBER_LABELS = {
    NumberClass.GUARD: "G",
    NumberClass.ONE: "1",
    NumberClass.TWO: "2",
    NumberClass.THREE: "3",
    NumberClass.FOUR: "4",
    NumberClass.FIVE: "5",
    NumberClass.OUTPOST: "O",
    NumberClass.BASE_SMALL: "Bs",
    NumberClass.BASE_LARGE: "Bb",
    NumberClass.NEGATIVE: "negative",
}


class ModelColor(Enum):
    RED = 0
    BLUE = 1
    GRAY = 2
    PURPLE = 3


COLOR_LABELS = {
    ModelColor.RED: "RED",
    ModelColor.BLUE: "BLUE",
    ModelColor.GRAY: "GRAY",
    ModelColor.PURPLE: "PURPLE",
}


@dataclass
class YoloObject:
    """YOLO detection result"""
    landmarks: np.ndarray  # 4 keypoints: [x1,y1,x2,y2,x3,y3,x4,y4] 左上逆时针
    label: int            # Number class (0-9)
    prob: float          # Confidence
    color_id: int        # Color class (0-3)


@dataclass
class Armor:
    """Armor with 6 landmarks for PnP"""
    landmarks_6: np.ndarray  # 6 points: [bottom_left, left_center, top_left, top_right, right_center, bottom_right]
    center: Tuple[float, float]
    armor_type: str  # "SMALL" or "LARGE"
    number_class: NumberClass
    confidence: float


def sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))


def apply_nms(boxes: List[np.ndarray], scores: List[float], iou_threshold: float = 0.45) -> List[int]:
    """Apply Non-Maximum Suppression"""
    if not boxes:
        return []

    # Sort by score
    indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)

    keep = []
    while indices:
        current = indices[0]
        keep.append(current)

        if len(indices) == 1:
            break

        # Calculate IoU with remaining boxes
        current_box = boxes[current]
        ious = []
        remaining = []
        for idx in indices[1:]:
            box = boxes[idx]
            # Calculate intersection
            x1 = max(current_box[0], box[0])
            y1 = max(current_box[1], box[1])
            x2 = min(current_box[2], box[2])
            y2 = min(current_box[3], box[3])

            if x2 < x1 or y2 < y1:
                ious.append(0.0)
            else:
                inter_area = (x2 - x1) * (y2 - y1)
                current_area = (current_box[2] - current_box[0]) * (current_box[3] - current_box[1])
                box_area = (box[2] - box[0]) * (box[3] - box[1])
                union_area = current_area + box_area - inter_area
                ious.append(inter_area / union_area if union_area > 0 else 0.0)

        # Keep boxes with IoU below threshold
        indices = [indices[i + 1] for i in range(len(indices) - 1) if ious[i] <= iou_threshold]

    return keep


class YoloDetector:
    def __init__(
        self,
        model_path: str,
        conf_threshold: float = 0.65,
        nms_threshold: float = 0.45,
        detect_color: int = 0,  # 0=RED, 1=BLUE
        use_coreml: bool = True,
    ):
        """
        Initialize YOLO detector

        Args:
            model_path: Path to ONNX model
            conf_threshold: Confidence threshold
            nms_threshold: NMS threshold
            detect_color: 0=RED, 1=BLUE
            use_coreml: Use CoreML for acceleration on Mac
        """
        self.conf_threshold = conf_threshold
        self.nms_threshold = nms_threshold
        self.detect_color = detect_color  # 0=red, 1=blue

        # Model parameters
        self.input_size = (640, 640)
        self.num_keypoints = 4
        self.num_color_classes = 4
        self.num_number_classes = 10  # G, 1, 2, 3, 4, 5, O, Bs, Bb, negative

        # Output shape: [25200, 23] - 4 kps + 1 conf + 4 color + 10 number
        self.output_rows = 25200
        self.output_cols = 23

        # Initialize ONNX Runtime
        sess_options = ort.SessionOptions()

        if use_coreml and sys.platform == 'darwin':
            # Use CoreML for Mac GPU acceleration (disable to avoid float16 type issues)
            providers = ['CPUExecutionProvider']
        else:
            providers = ['CPUExecutionProvider']

        print(f"Loading model from: {model_path}")
        print(f"Using providers: {providers}")

        self.session = ort.InferenceSession(model_path, sess_options, providers=providers)

        # Get input name
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name

        print(f"Input: {self.input_name}, Output: {self.output_name}")
        print("Model loaded successfully!")

    def preprocess(self, img: np.ndarray) -> np.ndarray:
        """Preprocess image for model input"""
        # Resize
        resized = cv2.resize(img, self.input_size)

        # BGR to RGB
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)

        # Normalize to [0, 1]
        normalized = rgb.astype(np.float32) / 255.0

        # HWC to CHW
        transposed = np.transpose(normalized, (2, 0, 1))

        # Add batch dimension
        batched = np.expand_dims(transposed, axis=0)

        # Convert to float16 for model input
        batched = batched.astype(np.float16)

        return batched

    def postprocess(self, output: np.ndarray) -> List[YoloObject]:
        """Postprocess model output"""
        objects = []

        # output shape: [1, 25200, 23] or [25200, 23]
        if output.ndim == 3:
            output = output[0]  # Remove batch dimension

        for i in range(self.output_rows):
            row = output[i]

            # Get confidence (index 8)
            conf = float(sigmoid(row[8]))
            if conf < self.conf_threshold:
                continue

            # Get color scores (indices 9-12)
            color_scores = sigmoid(row[9:13])
            color_id = int(np.argmax(color_scores))
            max_color_score = float(color_scores[color_id])

            # Filter by detect_color
            # 0=RED in detect_color means detect RED, model RED=0, BLUE=1
            if color_id in [2, 3]:  # GRAY or PURPLE - skip
                continue
            if self.detect_color == 0 and color_id == 1:  # Want RED but found BLUE
                continue
            if self.detect_color == 1 and color_id == 0:  # Want BLUE but found RED
                continue

            # Get number scores (indices 13-22, 10 classes)
            num_scores = sigmoid(row[13:23])
            num_id = int(np.argmax(num_scores))
            max_num_score = float(num_scores[num_id])

            # Create YoloObject
            obj = YoloObject(
                landmarks=row[0:8].astype(np.float32),
                label=num_id,
                prob=conf * max_color_score,
                color_id=color_id,
            )
            objects.append(obj)

        return objects

    def detect(self, img: np.ndarray) -> List[YoloObject]:
        """Detect armors in image"""
        # Preprocess
        input_data = self.preprocess(img)

        # Run inference
        outputs = self.session.run([self.output_name], {self.input_name: input_data})
        output = outputs[0]

        # Postprocess
        return self.postprocess(output)

    def draw_detection(self, img: np.ndarray, objects: List[YoloObject], draw_landmarks: bool = True) -> np.ndarray:
        """Draw detection results on image"""
        result = img.copy()

        # Calculate scale factors to map from 640x640 to original image size
        orig_h, orig_w = result.shape[:2]
        scale_x = orig_w / self.input_size[0]  # 640
        scale_y = orig_h / self.input_size[1]  # 640

        for obj in objects:
            # Scale landmarks from 640x640 to original image size
            landmarks = obj.landmarks.reshape(-1, 2)  # 4 points
            landmarks_scaled = landmarks * np.array([scale_x, scale_y])

            # Determine color for drawing
            if obj.color_id == 0:  # RED
                color = (0, 0, 255)
            elif obj.color_id == 1:  # BLUE
                color = (255, 0, 0)
            elif obj.color_id == 2:  # GRAY
                color = (128, 128, 128)
            else:  # PURPLE
                color = (255, 0, 255)

            if draw_landmarks:
                # Draw polygon from 4 keypoints
                pts = landmarks_scaled.astype(np.int32).reshape((-1, 1, 2))
                cv2.polylines(result, [pts], isClosed=True, color=color, thickness=2)

                # Draw center point
                center = np.mean(landmarks_scaled, axis=0).astype(np.int32)
                cv2.circle(result, tuple(center), 5, color, -1)

                # Draw label
                label = NUMBER_LABELS[NumberClass(obj.label)]
                text_color = (0, 255, 0)
                cv2.putText(result, f"{label}", tuple(center + np.array([10, -10])),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, text_color, 2)

                # Print keypoint coordinates for debugging
                for j, pt in enumerate(landmarks_scaled):
                    pt_int = tuple(pt.astype(np.int32))
                    cv2.putText(result, f"P{j}({pt[0]:.0f},{pt[1]:.0f})",
                               (pt_int[0], pt_int[1] - 10 - j*15),
                               cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 255), 1)
            else:
                # Draw bounding box
                x_coords = landmarks_scaled[:, 0]
                y_coords = landmarks_scaled[:, 1]
                x1, y1 = int(x_coords.min()), int(y_coords.min())
                x2, y2 = int(x_coords.max()), int(y_coords.max())
                cv2.rectangle(result, (x1, y1), (x2, y2), color, 2)

                label = NUMBER_LABELS[NumberClass(obj.label)]
                cv2.putText(result, f"{label}", (x1, y1 - 10),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

        return result


def main():
    parser = argparse.ArgumentParser(description='YOLO Armor Detection on Mac')
    parser.add_argument('--model', type=str,
                       default='/Users/chenyouyuan/cyy_ws/jetson_nx_yolo/RobotDetectionModel/Model/0526.onnx',
                       help='Path to ONNX model')
    parser.add_argument('--image', type=str, required=True,
                       help='Path to input image')
    parser.add_argument('--output', type=str, default='output.jpg',
                       help='Path to output image')
    parser.add_argument('--conf', type=float, default=0.65,
                       help='Confidence threshold')
    parser.add_argument('--nms', type=float, default=0.45,
                       help='NMS threshold')
    parser.add_argument('--color', type=int, default=0, choices=[0, 1],
                       help='Detection color: 0=RED, 1=BLUE')
    parser.add_argument('--no-coreml', action='store_true',
                       help='Disable CoreML acceleration')
    parser.add_argument('--print-raw', action='store_true',
                       help='Print raw model output for first detection')
    parser.add_argument('--video', type=str,
                       help='Use video instead of image')
    args = parser.parse_args()

    # Initialize detector
    detector = YoloDetector(
        model_path=args.model,
        conf_threshold=args.conf,
        nms_threshold=args.nms,
        detect_color=args.color,
        use_coreml=not args.no_coreml,
    )

    if args.video:
        # Video mode
        cap = cv2.VideoCapture(args.video)
        if not cap.isOpened():
            print(f"Error: Cannot open video {args.video}")
            return

        window_name = "YOLO Detection"
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            # Detect
            objects = detector.detect(frame)

            # Draw and display
            result = detector.draw_detection(frame, objects)
            cv2.imshow(window_name, result)

            # Print detection info
            print(f"Detected {len(objects)} objects")
            for i, obj in enumerate(objects):
                print(f"  [{i}] label={NUMBER_LABELS[NumberClass(obj.label)]}, "
                      f"color={COLOR_LABELS[ModelColor(obj.color_id)]}, "
                      f"conf={obj.prob:.3f}")
                print(f"       landmarks: {obj.landmarks}")

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q') or key == 27:
                break

        cap.release()
        cv2.destroyAllWindows()

    else:
        # Image mode
        img = cv2.imread(args.image)
        if img is None:
            print(f"Error: Cannot read image {args.image}")
            return

        # Detect
        start_time = time.time()
        objects = detector.detect(img)
        elapsed = time.time() - start_time

        # Draw
        result = detector.draw_detection(img, objects)

        # Save output
        cv2.imwrite(args.output, result)
        print(f"\nSaved result to {args.output}")

        # Print detection info
        print(f"\nInference time: {elapsed*1000:.2f} ms")
        print(f"Detected {len(objects)} objects")

        for i, obj in enumerate(objects):
            print(f"\n[{i}] {NUMBER_LABELS[NumberClass(obj.label)]} ({COLOR_LABELS[ModelColor(obj.color_id)]})")
            print(f"    Confidence: {obj.prob:.4f}")
            print(f"    Raw landmarks (4 keypoints, 左上逆时针):")
            print(f"      P0(top_left):      ({obj.landmarks[0]:.2f}, {obj.landmarks[1]:.2f})")
            print(f"      P1(top_right):     ({obj.landmarks[2]:.2f}, {obj.landmarks[3]:.2f})")
            print(f"      P2(bottom_right):  ({obj.landmarks[4]:.2f}, {obj.landmarks[5]:.2f})")
            print(f"      P3(bottom_left):   ({obj.landmarks[6]:.2f}, {obj.landmarks[7]:.2f})")

            # Calculate center and bounding box
            pts = obj.landmarks.reshape(-1, 2)
            center = np.mean(pts, axis=0)
            print(f"    Center (mean): ({center[0]:.2f}, {center[1]:.2f})")
            print(f"    BBox: x=[{pts[:,0].min():.0f}, {pts[:,0].max():.0f}], y=[{pts[:,1].min():.0f}, {pts[:,1].max():.0f}]")

        # Display
        cv2.imshow("Detection Result", result)
        print("\nPress any key to exit...")
        cv2.waitKey(0)
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()
