# Detection logic using OpenVINO models (YOLO, etc.)

import os
import sys
import time
import cv2
import numpy as np
from pathlib import Path
from typing import List, Dict, Tuple, Optional

# --- Install required packages if missing ---
try:
    import openvino as ov
except ImportError:
    print("Installing openvino...")
    os.system('pip install --quiet "openvino>=2024.0.0"')
    import openvino as ov
try:
    from ultralytics import YOLO
except ImportError:
    print("Installing ultralytics...")
    os.system('pip install --quiet "ultralytics==8.3.0"')
    from ultralytics import YOLO
try:
    import nncf
except ImportError:
    print("Installing nncf...")
    os.system('pip install --quiet "nncf>=2.9.0"')
    import nncf

# --- COCO dataset class names ---
COCO_CLASSES = {
    0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 
    6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 
    11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 
    16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 
    22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 
    27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 
    32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 
    36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 
    40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 
    46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 
    51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 
    57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 
    62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 
    67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 
    72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 
    77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'
}

# Traffic-related classes we're interested in (using standard COCO indices)
TRAFFIC_CLASS_NAMES = [
    'person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck',
    'traffic light', 'stop sign', 'parking meter'
]

# --- Model Conversion and Quantization ---
def convert_yolo_to_openvino(model_name: str = "yolo11x", half: bool = True) -> Path:
    """Convert YOLOv11x PyTorch model to OpenVINO IR format."""
    pt_path = Path(f"{model_name}.pt")
    ov_dir = Path(f"{model_name}_openvino_model")
    ov_xml = ov_dir / f"{model_name}.xml"
    if not ov_xml.exists():
        print(f"Exporting {pt_path} to OpenVINO IR...")
        model = YOLO(str(pt_path))
        model.export(format="openvino", dynamic=True, half=half)
    else:
        print(f"OpenVINO IR already exists: {ov_xml}")
    return ov_xml

def quantize_openvino_model(ov_xml: Path, model_name: str = "yolo11x") -> Path:
    """Quantize OpenVINO IR model to INT8 using NNCF."""
    int8_dir = Path(f"{model_name}_openvino_int8_model")
    int8_xml = int8_dir / f"{model_name}.xml"
    if int8_xml.exists():
        print(f"INT8 model already exists: {int8_xml}")
        return int8_xml
    print("Quantization requires a calibration dataset. Skipping actual quantization in this demo.")
    return ov_xml  # Return FP32 if no quantization

class OpenVINOVehicleDetector:
    def __init__(self, model_path: str = None, device: str = "AUTO", use_quantized: bool = False, enable_ocr: bool = False, confidence_threshold: float = 0.4):
        import openvino as ov
        self.device = device
        self.confidence_threshold = confidence_threshold
        self.ocr_reader = None
        self.class_names = TRAFFIC_CLASS_NAMES
        self.performance_stats = {
            'fps': 0,
            'avg_inference_time': 0,
            'frames_processed': 0,
            'backend': f"OpenVINO-{device}",
            'total_detections': 0,
            'detection_rate': 0
        }
        self._inference_times = []
        self._start_time = time.time()
        self._frame_count = 0
        
        # Model selection logic
        self.model_path = self._find_best_model(model_path, use_quantized)
        print(f"🎯 OpenVINOVehicleDetector: Using model: {self.model_path}")
        
        self.core = ov.Core()
        self.model = self.core.read_model(self.model_path)
        # Always reshape to static shape before accessing .shape
        self.model.reshape({0: [1, 3, 640, 640]})
        self.input_shape = self.model.inputs[0].shape
        self.input_height = self.input_shape[2]
        self.input_width = self.input_shape[3]
        self.ov_config = {}
        if device != "CPU":
            # Already reshaped above, so nothing more needed here
            pass
        if "GPU" in device or ("AUTO" in device and "GPU" in self.core.available_devices):
            self.ov_config = {"GPU_DISABLE_WINOGRAD_CONVOLUTION": "YES"}
        self.compiled_model = self.core.compile_model(model=self.model, device_name=self.device, config=self.ov_config)

        self.output_layer = self.compiled_model.output(0)

    def _find_best_model(self, model_path, use_quantized):
        # If a specific model path is provided, use it directly
        if model_path and Path(model_path).exists():
            print(f"🎯 Using provided model path: {model_path}")
            return str(model_path)
        
        # If no model path provided, extract model name from path or default to yolo11x
        model_name = "yolo11x"  # Default fallback
        if model_path:
            # Try to extract model name from path
            path_obj = Path(model_path)
            if "yolo11n" in str(path_obj).lower():
                model_name = "yolo11n"
            elif "yolo11s" in str(path_obj).lower():
                model_name = "yolo11s"
            elif "yolo11m" in str(path_obj).lower():
                model_name = "yolo11m"
            elif "yolo11l" in str(path_obj).lower():
                model_name = "yolo11l"
            elif "yolo11x" in str(path_obj).lower():
                model_name = "yolo11x"
        
        print(f"🔍 Searching for {model_name} model files...")
        
        # Priority: quantized IR > IR > .pt
        search_paths = [
            Path(f"{model_name}_openvino_int8_model/{model_name}.xml") if use_quantized else None,
            Path(f"{model_name}_openvino_model/{model_name}.xml"),
            Path(f"rcb/{model_name}_openvino_model/{model_name}.xml"),
            Path(f"{model_name}.xml"),
            Path(f"rcb/{model_name}.xml"),
            Path(f"{model_name}.pt"),
            Path(f"rcb/{model_name}.pt")
        ]
        
        for p in search_paths:
            if p and p.exists():
                print(f"✅ Found model: {p}")
                return str(p)
        
        # Fallback to any yolo11x if specific model not found
        fallback_paths = [
            Path("yolo11x_openvino_model/yolo11x.xml"),
            Path("rcb/yolo11x_openvino_model/yolo11x.xml"),
            Path("yolo11x.xml"),
            Path("rcb/yolo11x.xml"),
            Path("yolo11x.pt"),
            Path("rcb/yolo11x.pt")
        ]
        
        for p in fallback_paths:
            if p and p.exists():
                print(f"⚠️ Using fallback model: {p}")
                return str(p)
                
        raise FileNotFoundError(f"No suitable {model_name} model found for OpenVINO.")

    def detect_vehicles(self, frame: np.ndarray, conf_threshold: float = None) -> List[Dict]:
        if conf_threshold is None:
            conf_threshold = 0.1  # Lowered for debugging
        start = time.time()
        input_tensor = self._preprocess(frame)
        output = self.compiled_model([input_tensor])[self.output_layer]
        # Debug: print raw output shape
        # print(f"[DEBUG] Model output shape: {output.shape}")
        detections = self._postprocess(output, frame.shape, conf_threshold)
        # print(f"[DEBUG] Detections after postprocess: {len(detections)}")
        elapsed = time.time() - start
        self._inference_times.append(elapsed)
        self._frame_count += 1
        self.performance_stats['frames_processed'] = self._frame_count
        self.performance_stats['total_detections'] += len(detections)
        if len(self._inference_times) > 100:
            self._inference_times.pop(0)
        self.performance_stats['avg_inference_time'] = float(np.mean(self._inference_times)) if self._inference_times else 0
        total_time = time.time() - self._start_time
        self.performance_stats['fps'] = self._frame_count / total_time if total_time > 0 else 0
        return detections

    def _preprocess(self, frame: np.ndarray) -> np.ndarray:
        img = cv2.resize(frame, (self.input_width, self.input_height))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.0
        img = img.transpose(2, 0, 1)[None]
        return img

    def _postprocess(self, output: np.ndarray, frame_shape, conf_threshold: float) -> List[Dict]:
        # Output: (1, 84, 8400) or (84, 8400) or (8400, 84)
        if output.ndim == 3:
            output = np.squeeze(output)
        if output.shape[0] == 84:
            output = output.T  # (8400, 84)
        boxes = output[:, :4]
        scores = output[:, 4:]
        class_ids = np.argmax(scores, axis=1)
        confidences = np.max(scores, axis=1)
        detections = []
        h, w = frame_shape[:2]
        for i, (box, score, class_id) in enumerate(zip(boxes, confidences, class_ids)):
            if score < conf_threshold:
                continue
            x_c, y_c, bw, bh = box
            # If normalized, scale to input size
            if all(0.0 <= v <= 1.0 for v in box):
                x_c *= self.input_width
                y_c *= self.input_height
                bw *= self.input_width
                bh *= self.input_height
            # Scale to original frame size
            scale_x = w / self.input_width
            scale_y = h / self.input_height
            x_c *= scale_x
            y_c *= scale_y
            bw *= scale_x
            bh *= scale_y
            x1 = int(round(x_c - bw / 2))
            y1 = int(round(y_c - bh / 2))
            x2 = int(round(x_c + bw / 2))
            y2 = int(round(y_c + bh / 2))
            x1 = max(0, min(x1, w - 1))
            y1 = max(0, min(y1, h - 1))
            x2 = max(0, min(x2, w - 1))
            y2 = max(0, min(y2, h - 1))
            if x2 <= x1 or y2 <= y1:
                continue
            # Only keep class 9 as traffic light, rename if found
            if class_id == 9:
                class_name = "traffic light"
            elif class_id < len(TRAFFIC_CLASS_NAMES):
                class_name = TRAFFIC_CLASS_NAMES[class_id]
            else:
                continue  # Remove unknown/other classes
            detections.append({
                'bbox': [x1, y1, x2, y2],
                'confidence': float(score),
                'class_id': int(class_id),
                'class_name': class_name
            })
        # print(f"[DEBUG] Raw detections before NMS: {len(detections)}")
        # Apply NMS
        if len(detections) > 0:
            boxes = np.array([det['bbox'] for det in detections])
            scores = np.array([det['confidence'] for det in detections])
            indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_threshold, 0.5)
            if isinstance(indices, (list, tuple)) and len(indices) > 0:
                indices = np.array(indices).flatten()
            elif isinstance(indices, np.ndarray) and indices.size > 0:
                indices = indices.flatten()
            else:
                indices = []
            detections = [detections[int(i)] for i in indices] if len(indices) > 0 else []
        # print(f"[DEBUG] Detections after NMS: {len(detections)}")
        return detections

    def draw(self, frame: np.ndarray, detections: List[Dict], box_thickness: int = 2) -> np.ndarray:
        # 80+ visually distinct colors for COCO classes (BGR)
        COCO_COLORS = [
            (255, 56, 56), (255, 157, 151), (255, 112, 31), (255, 178, 29), (207, 210, 49),
            (72, 249, 10), (146, 204, 23), (61, 219, 134), (26, 147, 52), (0, 212, 187),
            (44, 153, 168), (0, 194, 255), (52, 69, 147), (100, 115, 255), (0, 24, 236),
            (132, 56, 255), (82, 0, 133), (203, 56, 255), (255, 149, 200), (255, 55, 199),
            (255, 255, 56), (255, 255, 151), (255, 255, 31), (255, 255, 29), (207, 255, 49),
            (72, 255, 10), (146, 255, 23), (61, 255, 134), (26, 255, 52), (0, 255, 187),
            (44, 255, 168), (0, 255, 255), (52, 255, 147), (100, 255, 255), (0, 255, 236),
            (132, 255, 255), (82, 255, 133), (203, 255, 255), (255, 255, 200), (255, 255, 199),
            (56, 255, 255), (157, 255, 151), (112, 255, 31), (178, 255, 29), (210, 255, 49),
            (249, 255, 10), (204, 255, 23), (219, 255, 134), (147, 255, 52), (212, 255, 187),
            (153, 255, 168), (194, 255, 255), (69, 255, 147), (115, 255, 255), (24, 255, 236),
            (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49),
            (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187),
            (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236),
            (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49),
            (249, 72, 10), (204, 146, 23), (219, 61, 134), (147, 26, 52), (212, 0, 187),
            (153, 44, 168), (194, 0, 255), (69, 52, 147), (115, 100, 255), (24, 0, 236),
            (56, 132, 255), (157, 82, 151), (112, 203, 31), (178, 255, 29), (210, 255, 49)
        ]
        for det in detections:
            x1, y1, x2, y2 = det['bbox']
            label = f"{det['class_name']} {det['confidence']:.2f}"
            color = COCO_COLORS[det['class_id'] % len(COCO_COLORS)]
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, box_thickness)
            cv2.putText(frame, label, (x1, max(y1 - 10, 0)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
        return frame
    
    def get_device(self):
        """Get the device being used for inference"""
        return self.device

if __name__ == "__main__":
    # Test the detector with YOLOv11n model
    detector = OpenVINOVehicleDetector(model_path="yolo11n_openvino_model/yolo11n.xml")
    print(f"Detector initialized with model: {detector.model_path}")
