import os
import cv2
import numpy as np
import onnxruntime as ort

from data_generator import CLASSES

class FruitVegetableDetector:
    def __init__(self, model_path=os.path.join("models", "model.onnx")):
        self.model_path = model_path
        self.session = None
        self.input_name = None
        self.load_model()
        
    def load_model(self):
        if not os.path.exists(self.model_path):
            print(f"Model file {self.model_path} not found. Running in SIMULATED mode.")
            self.session = None
            return
            
        try:
            # Set providers. SSDLite ONNX Runtime runs extremely fast on CPU
            providers = ['CPUExecutionProvider']
            if 'DirectMLExecutionProvider' in ort.get_available_providers():
                providers = ['DirectMLExecutionProvider', 'CPUExecutionProvider']
                
            self.session = ort.InferenceSession(self.model_path, providers=providers)
            self.input_name = self.session.get_inputs()[0].name
            print(f"ONNX model loaded successfully from {self.model_path}")
        except Exception as e:
            print(f"Failed to load ONNX model: {e}. Falling back to SIMULATED mode.")
            self.session = None

    def detect(self, image_rgb, confidence_threshold=0.85):
        """
        Runs inference on an RGB image.
        Args:
            image_rgb (np.ndarray): HWC RGB image.
            confidence_threshold (float): Minimum confidence to report.
        Returns:
            list: List of predictions with class, confidence, and box [xmin, ymin, xmax, ymax]
        """
        orig_h, orig_w, _ = image_rgb.shape
        
        # If no model is loaded, run simulation fallback
        if self.session is None:
            return self._simulate_detection(orig_w, orig_h, confidence_threshold)
            
        # Preprocessing: Resize to 320x320
        resized = cv2.resize(image_rgb, (320, 320))
        # Scale to [0.0, 1.0] and transpose to [C, H, W]
        input_tensor = resized.astype(np.float32) / 255.0
        input_tensor = np.transpose(input_tensor, (2, 0, 1)) # shape: (3, 320, 320)
        
        # Run inference
        outputs = self.session.run(None, {self.input_name: input_tensor})
        boxes, scores, labels = outputs
        
        predictions = []
        for i in range(len(scores)):
            score = float(scores[i])
            if score < confidence_threshold:
                continue
                
            # Class ID is 1-indexed, convert to 0-indexed
            class_id = int(labels[i]) - 1
            if class_id < 0 or class_id >= len(CLASSES):
                continue
                
            class_name = CLASSES[class_id]
            box_320 = boxes[i]
            
            # Scale box back to original coordinates
            xmin = int(round((box_320[0] / 320.0) * orig_w))
            ymin = int(round((box_320[1] / 320.0) * orig_h))
            xmax = int(round((box_320[2] / 320.0) * orig_w))
            ymax = int(round((box_320[3] / 320.0) * orig_h))
            
            # Clip to image boundaries
            xmin = max(0, min(xmin, orig_w - 1))
            ymin = max(0, min(ymin, orig_h - 1))
            xmax = max(xmin + 1, min(xmax, orig_w))
            ymax = max(ymin + 1, min(ymax, orig_h))
            
            predictions.append({
                "class": class_name,
                "confidence": round(score, 3),
                "box": [xmin, ymin, xmax, ymax]
            })
            
        # Sort predictions by confidence descending
        predictions = sorted(predictions, key=lambda x: x["confidence"], reverse=True)
        
        # Apply Non-Maximum Suppression (NMS) to eliminate duplicate/overlapping boxes
        predictions = self.apply_nms(predictions, iou_threshold=0.5)
        
        return predictions

    def apply_nms(self, predictions, iou_threshold=0.5):
        """
        Applies Non-Maximum Suppression (NMS) to predictions.
        """
        if not predictions:
            return []
            
        keep = []
        # Predictions are already sorted by confidence descending
        candidates = list(predictions)
        
        while candidates:
            best = candidates.pop(0)
            keep.append(best)
            
            remaining = []
            for pred in candidates:
                if self.calculate_iou(best["box"], pred["box"]) < iou_threshold:
                    remaining.append(pred)
            candidates = remaining
            
        return keep

    def calculate_iou(self, box1, box2):
        """
        Calculates Intersection over Union (IoU) of two boxes [xmin, ymin, xmax, ymax].
        """
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])
        
        intersection = max(0, x2 - x1) * max(0, y2 - y1)
        
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        union = area1 + area2 - intersection
        
        if union == 0:
            return 0
        return intersection / union

    def _simulate_detection(self, width, height, confidence_threshold):
        """
        Generates simulated detections based on random chances.
        Used as a fallback when the model hasn't been trained yet.
        """
        # Seed a pseudo-random detection based on image size to look semi-stable
        seed_factor = (width + height) % 6
        
        # Simulate 1 to 2 detections
        predictions = []
        
        # 1. First main prediction (high confidence)
        class_idx = seed_factor
        conf = 0.85 + (seed_factor * 0.02)
        
        # Simulate reasonable bounding box centered on scale
        cx, cy = int(width * 0.5), int(height * 0.5)
        w, h = int(width * 0.35), int(height * 0.45)
        
        predictions.append({
            "class": CLASSES[class_idx],
            "confidence": round(conf, 3),
            "box": [cx - w//2, cy - h//2, cx + w//2, cy + h//2]
        })
        
        # 2. Second prediction (low confidence, e.g. alternative suggestion)
        alt_class_idx = (class_idx + 1) % len(CLASSES)
        alt_conf = 0.12 + (seed_factor * 0.01)
        
        if alt_conf >= confidence_threshold:
            predictions.append({
                "class": CLASSES[alt_class_idx],
                "confidence": round(alt_conf, 3),
                "box": [cx - w//2 + 20, cy - h//2 + 10, cx + w//2 + 20, cy + h//2 + 10]
            })
            
        return predictions