from models.predictions import PredictionCollection
from models.cards import Card
from collections import Counter, deque
from ultralytics import YOLO
from cv2.typing import MatLike

value_model = YOLO("ai_models/value.pt")
symbol_model = YOLO("ai_models/symbol.pt")

def predict_with_model(frame: MatLike, model: YOLO):
    results = model(frame, imgsz=640, verbose=False)
    return PredictionCollection(results[0])


def calculate_iou(box1: list[float], box2: list[float]) -> float:
    """
    Calculate Intersection over Union (IoU) for normalized boxes
    :param box: [x1, y1, x2, y2] in normalized coordinates (0-1)
    """
    # Convert to coordinates
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    
    # Calculate area of intersection
    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    
    # Calculate area of union
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area1 + area2 - intersection
    
    return intersection / union if union > 0 else 0

class PredictService:

    @staticmethod
    def predict_frame(frame: MatLike):
        return {
            "symbol" : predict_with_model(frame, symbol_model),
            "value" : predict_with_model(frame, value_model)
        }
    
    @staticmethod
    def pair_predictions(full_prediction: dict[str, PredictionCollection], iou_threshold: float = 0.7) -> list[Card]:
        """
        Pair symbol and value predictions based on bounding box similarity
        """
        pairs = []
        used_values = set()
        
        symbols = full_prediction["symbol"].predictions
        values = full_prediction["value"].predictions
        
        for symbol_pred in symbols:
            best_iou = -1
            best_value = None
            
            for idx, value_pred in enumerate(values):
                if idx in used_values:
                    continue
                
                iou = calculate_iou(symbol_pred.xyxyn, value_pred.xyxyn)
                if iou > best_iou and iou >= iou_threshold:
                    best_iou = iou
                    best_value = idx
            
            if best_value is not None:
                used_values.add(best_value)
                pairs.append(Card(
                    symbol_pred.cls,
                    values[best_value].cls
                ))
        
        return pairs

class PredictionSmoother:
    def __init__(self, maxlen: int):
        self.prediction_queue = deque(maxlen=maxlen)
        self.current_values: tuple[Card] = ()

    def add_prediction(self, cards: list[Card]):
        sorted_cards_tuple = tuple(sorted(cards, key=lambda c: (c.symbol, c.value)))
        self.prediction_queue.append(sorted_cards_tuple)
        most_common_tuple, _ = Counter(self.prediction_queue).most_common(1)[0]
        if most_common_tuple == self.current_values:
            return None
        else:
            self.current_values = most_common_tuple
            return list(most_common_tuple)