import cv2
import numpy as np
from sklearn.cluster import KMeans
import football_analytics.config as config

class TeamClassifier:
    def __init__(self, n_clusters=2, distance_threshold=45.0):
        """
        Unsupervised jersey color classifier using KMeans.
        
        Args:
            n_clusters (int): Number of teams to cluster (default 2)
            distance_threshold (float): Euclidean distance threshold in RGB space (0-255)
                                         above which a player is classified as 'Other' (referee/goalkeeper).
        """
        self.n_clusters = n_clusters
        self.distance_threshold = distance_threshold
        self.kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
        self.is_fitted = False
        self.team_colors = {}  # Store cluster centers for visual annotation
        self.features_buffer = []

    def _extract_jersey_color(self, frame, bbox):
        """
        Crops the jersey area and returns the average RGB color.
        """
        h_frame, w_frame, _ = frame.shape
        x1, y1, x2, y2 = bbox
        
        # Clip to frame boundaries
        x1, y1 = max(0, int(x1)), max(0, int(y1))
        x2, y2 = min(w_frame, int(x2)), min(h_frame, int(y2))
        
        w = x2 - x1
        h = y2 - y1
        
        if w <= 0 or h <= 0:
            return None
            
        # Get chest region based on config ratios
        ymin_j = int(y1 + h * config.JERSEY_BOX[0])
        ymax_j = int(y1 + h * config.JERSEY_BOX[2])
        xmin_j = int(x1 + w * config.JERSEY_BOX[1])
        xmax_j = int(x1 + w * config.JERSEY_BOX[3])
        
        crop = frame[ymin_j:ymax_j, xmin_j:xmax_j]
        if crop.size == 0:
            return None
            
        # Convert BGR to RGB
        crop_rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
        
        # Calculate mean RGB color of the chest region
        mean_color = crop_rgb.mean(axis=(0, 1))
        return mean_color

    def collect_features(self, frame, detections):
        """
        Extracts jersey color features for all player detections in a frame and adds to buffer.
        """
        for bbox, class_id in zip(detections.xyxy, detections.class_id):
            if class_id == 0:  # Person (Player/Referee)
                color = self._extract_jersey_color(frame, bbox)
                if color is not None:
                    self.features_buffer.append(color)

    def fit(self):
        """
        Fits the KMeans model on the accumulated color features.
        """
        if getattr(config, "USE_SUPERVISED_TEAMS", False):
            print("[Team Classifier] Supervised team mode enabled. Using predefined target colors:")
            self.team_colors = {}
            for k, v in config.SUPERVISED_COLORS.items():
                print(f"  Team {k} RGB: {v}")
                self.team_colors[k] = np.array(v)
            self.is_fitted = True
            return True

        if len(self.features_buffer) < self.n_clusters * 10:
            print(f"[Team Classifier] Warning: Not enough player detections to fit KMeans ({len(self.features_buffer)} found). Using fallback colors.")
            # Create a mock fit or use default centers
            self.team_colors = {0: np.array([255, 0, 0]), 1: np.array([0, 0, 255])}
            self.is_fitted = True
            return False

        features = np.array(self.features_buffer)
        self.kmeans.fit(features)
        self.is_fitted = True
        
        # Store team average colors
        for i in range(self.n_clusters):
            self.team_colors[i] = self.kmeans.cluster_centers_[i]
            
        print(f"[Team Classifier] Successfully fitted K-Means on {len(features)} detections.")
        print(f"Team 0 Center Color (RGB): {self.team_colors[0].astype(int)}")
        print(f"Team 1 Center Color (RGB): {self.team_colors[1].astype(int)}")
        return True

    def predict_team(self, frame, bbox):
        """
        Predicts the team of a player based on jersey color.
        
        Returns:
            int: 0 for Team A, 1 for Team B, 2 for Referee/Other
        """
        if not self.is_fitted:
            return 0  # Default to Team 0 if not fitted
            
        color = self._extract_jersey_color(frame, bbox)
        if color is None:
            return 0
            
        if getattr(config, "USE_SUPERVISED_TEAMS", False):
            # Supervised classification: find closest of the pre-defined target colors
            dists = {
                k: np.linalg.norm(color - np.array(v))
                for k, v in config.SUPERVISED_COLORS.items()
            }
            return min(dists, key=dists.get)

        # Calculate distances to all K-Means cluster centers
        distances = [np.linalg.norm(color - center) for center in self.kmeans.cluster_centers_]
        min_dist = min(distances)
        predicted_cluster = np.argmin(distances)
        
        # If distance to nearest team is too large, classify as referee/other (class ID 2)
        if min_dist > self.distance_threshold:
            return 2
            
        return int(predicted_cluster)