modules/visualizer.py

import os
import cv2
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # Headless backend for matplotlib
import matplotlib.pyplot as plt
from mplsoccer import Pitch
import football_analytics.config as config

class FootballVisualizer:
    def __init__(self):
        # Color palettes (BGR for OpenCV)
        self.colors = {
            0: (230, 50, 50),     # Team A: Bright Red/Coral
            1: (50, 100, 240),    # Team B: Bright Blue
            2: (50, 240, 240),    # Referee/Other: Yellow
            -1: (200, 200, 200),  # Unknown: Light Gray
            "ball": (0, 255, 255), # Ball: Neon Yellow/Green
            "radar_bg": (34, 49, 43), # Dark Green Pitch Background
            "radar_line": (200, 215, 205), # Pitch Lines
            "possession_glow": (0, 255, 0) # Glowing green for player in possession
        }
        
        # Dimensions of the radar overlay
        self.radar_w = 320
        self.radar_h = 200
        self.radar_margin = 20

    def draw_pitch_radar(self, players, ball_pos):
        """
        Draws a 2D mini-pitch radar representation.
        
        Args:
            players (list of dict): [{id, team, x_pitch, y_pitch}]
            ball_pos (tuple): (x_pitch, y_pitch) or None
            
        Returns:
            np.ndarray: Mini-pitch image of size (radar_h, radar_w, 3)
        """
        # Create dark green pitch background
        radar = np.zeros((self.radar_h, self.radar_w, 3), dtype=np.uint8)
        radar[:] = self.colors["radar_bg"]
        
        # Scale factors
        scale_x = self.radar_w / config.PITCH_WIDTH
        scale_y = self.radar_h / config.PITCH_HEIGHT
        
        def to_radar_coords(x, y):
            rx = int(x * scale_x)
            ry = int(y * scale_y)
            return rx, ry

        # Draw pitch boundaries
        cv2.rectangle(radar, (0, 0), (self.radar_w - 1, self.radar_h - 1), self.colors["radar_line"], 2)
        
        # Draw half-way line
        mid_x = int(self.radar_w / 2)
        cv2.line(radar, (mid_x, 0), (mid_x, self.radar_h), self.colors["radar_line"], 1)
        
        # Draw center circle
        center_circle_r = int(9.15 * scale_x)
        cv2.circle(radar, (mid_x, int(self.radar_h / 2)), center_circle_r, self.colors["radar_line"], 1)
        cv2.circle(radar, (mid_x, int(self.radar_h / 2)), 2, self.colors["radar_line"], -1)
        
        # Draw penalty areas (UEFA rules: 16.5m deep, 40.32m wide)
        # Left penalty area
        pa_w = int(16.5 * scale_x)
        pa_h = int(40.32 * scale_y)
        pa_y1 = int((config.PITCH_HEIGHT - 40.32) / 2 * scale_y)
        pa_y2 = pa_y1 + pa_h
        cv2.rectangle(radar, (0, pa_y1), (pa_w, pa_y2), self.colors["radar_line"], 1)
        # Goal area left
        ga_w = int(5.5 * scale_x)
        ga_h = int(18.32 * scale_y)
        ga_y1 = int((config.PITCH_HEIGHT - 18.32) / 2 * scale_y)
        ga_y2 = ga_y1 + ga_h
        cv2.rectangle(radar, (0, ga_y1), (ga_w, ga_y2), self.colors["radar_line"], 1)
        
        # Right penalty area
        cv2.rectangle(radar, (self.radar_w - pa_w, pa_y1), (self.radar_w, pa_y2), self.colors["radar_line"], 1)
        # Goal area right
        cv2.rectangle(radar, (self.radar_w - ga_w, ga_y1), (self.radar_w, ga_y2), self.colors["radar_line"], 1)

        # Draw players on radar
        for p in players:
            team = p["team"]
            px, py = p["x_pitch"], p["y_pitch"]
            rx, ry = to_radar_coords(px, py)
            
            # Skip if invalid coords
            if not (0 <= rx < self.radar_w and 0 <= ry < self.radar_h):
                continue
                
            color = self.colors.get(team, self.colors[-1])
            # Draw player dot
            cv2.circle(radar, (rx, ry), 5, color, -1)
            cv2.circle(radar, (rx, ry), 6, (255, 255, 255), 1)

        # Draw ball on radar
        if ball_pos is not None:
            bx, by = ball_pos
            rx, ry = to_radar_coords(bx, by)
            if 0 <= rx < self.radar_w and 0 <= ry < self.radar_h:
                cv2.circle(radar, (rx, ry), 4, self.colors["ball"], -1)
                cv2.circle(radar, (rx, ry), 5, (0, 0, 0), 1)
                
        return radar

    def annotate_frame(self, frame, detections, team_assignments, ball_xyxy, current_possession_player, draw_possession=True):
        """
        Annotates the frame with bounding boxes, team color coding, ball info, and overlays the 2D radar.
        
        Args:
            frame (np.ndarray): Input image frame.
            detections (sv.Detections): Player tracked detections.
            team_assignments (dict): track_id -> team_id.
            ball_xyxy (np.ndarray): Ball bounding box [x1, y1, x2, y2] or None.
            current_possession_player (int): ID of player currently in possession.
            
        Returns:
            np.ndarray: Annotated frame.
        """
        annotated_frame = frame.copy()
        h_frame, w_frame, _ = frame.shape
        
        # 1. Annotate players
        for i, bbox in enumerate(detections.xyxy):
            track_id = int(detections.tracker_id[i]) if detections.tracker_id is not None else -1
            class_id = int(detections.class_id[i])
            
            if class_id != 0:  # Skip ball (handle separately)
                continue
                
            team = team_assignments.get(track_id, -1)
            if team == 3:  # Skip drawing staff/coaches completely
                continue
            color = self.colors.get(team, self.colors[-1])
            
            x1, y1, x2, y2 = map(int, bbox)
            
            # Draw semi-transparent ellipse at the feet of the player for a premium tactical look
            feet_y = y2
            center_x = int((x1 + x2) / 2)
            axis_x = int((x2 - x1) / 2)
            axis_y = int(axis_x / 3)  # Perspective flattening
            
            # Make a glowing ellipse for possession
            is_possessor = (track_id == current_possession_player and track_id != -1)
            ellipse_color = self.colors["possession_glow"] if is_possessor else color
            thickness = 3 if is_possessor else 2
            
            # Draw ellipse feet base
            cv2.ellipse(annotated_frame, (center_x, feet_y), (axis_x, axis_y), 0, 0, 360, ellipse_color, thickness)
            
            # Draw bounding box (subtle corner-only or thin lines)
            cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 1)
            
            # Draw label background
            label_text = f"ID:{track_id}"
            if team == 0:
                label_text = f"A-{track_id}"
            elif team == 1:
                label_text = f"B-{track_id}"
            elif team == 2:
                label_text = f"REF"
                
            (w, h_text), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.4, 1)
            cv2.rectangle(annotated_frame, (x1, y1 - h_text - 4), (x1 + w + 4, y1), color, -1)
            cv2.putText(annotated_frame, label_text, (x1 + 2, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv2.LINE_AA)

        # 2. Annotate Ball
        if ball_xyxy is not None and len(ball_xyxy) > 0:
            if ball_xyxy.ndim == 2:
                ball_box = ball_xyxy[0]
            else:
                ball_box = ball_xyxy
                
            bx1, by1, bx2, by2 = map(int, ball_box)
            # Draw glowing circle around the ball
            bc_x = int((bx1 + bx2) / 2)
            bc_y = int((by1 + by2) / 2)
            r = max(6, int((bx2 - bx1) / 2))
            cv2.circle(annotated_frame, (bc_x, bc_y), r + 2, (0, 0, 0), 1)
            cv2.circle(annotated_frame, (bc_x, bc_y), r, self.colors["ball"], 2)
            
            # Draw a pointer above the ball
            pointer_y = by1 - 8
            pointer_pts = np.array([[bc_x, by1 - 2], [bc_x - 4, pointer_y], [bc_x + 4, pointer_y]], dtype=np.int32)
            cv2.fillPoly(annotated_frame, [pointer_pts], self.colors["ball"])

        # 3. Add possession banner overlay at the top middle
        if draw_possession:
            banner_w, banner_h = 360, 40
            banner_x = int((w_frame - banner_w) / 2)
            banner_y = 15
            
            # Transparent background for banner
            overlay = annotated_frame.copy()
            cv2.rectangle(overlay, (banner_x, banner_y), (banner_x + banner_w, banner_y + banner_h), (25, 25, 25), -1)
            cv2.addWeighted(overlay, 0.6, annotated_frame, 0.4, 0, annotated_frame)
            
            # Add possession text
            poss_text = "POSSESSION: CONTESTED"
            text_color = (200, 200, 200)
            
            if current_possession_player is not None:
                team_poss = team_assignments.get(current_possession_player, -1)
                if team_poss == 0:
                    poss_text = f"POSSESSION: TEAM A (ID {current_possession_player})"
                    text_color = (self.colors[0][2], self.colors[0][1], self.colors[0][0])
                elif team_poss == 1:
                    poss_text = f"POSSESSION: TEAM B (ID {current_possession_player})"
                    text_color = (self.colors[1][2], self.colors[1][1], self.colors[1][0])
                    
            (w_txt, h_txt), _ = cv2.getTextSize(poss_text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
            tx = banner_x + int((banner_w - w_txt) / 2)
            ty = banner_y + int((banner_h + h_txt) / 2)
            
            cv2.putText(annotated_frame, poss_text, (tx, ty), cv2.FONT_HERSHEY_SIMPLEX, 0.5, text_color, 2, cv2.LINE_AA)

        return annotated_frame

    def overlay_radar_on_frame(self, frame, radar_img):
        """
        Overlays the radar image in the top right corner with semi-transparency.
        """
        h_frame, w_frame, _ = frame.shape
        
        # Position: Top-right corner with 20px padding
        x_offset = w_frame - self.radar_w - self.radar_margin
        y_offset = self.radar_margin
        
        # Extract Region of Interest (ROI) from frame
        roi = frame[y_offset:y_offset + self.radar_h, x_offset:x_offset + self.radar_w]
        
        # Blend ROI and radar image (0.7 radar + 0.3 frame background for transparency)
        blended = cv2.addWeighted(radar_img, 0.75, roi, 0.25, 0)
        
        # Put blended radar back onto frame
        frame[y_offset:y_offset + self.radar_h, x_offset:x_offset + self.radar_w] = blended
        
        # Draw simple border around radar
        cv2.rectangle(frame, (x_offset, y_offset), (x_offset + self.radar_w, y_offset + self.radar_h), (120, 130, 125), 2)
        
        return frame

    def save_team_heatmaps(self, positions_log, output_dir=None):
        """
        Creates and saves kernel density heatmaps for each team.
        """
        out_dir = output_dir or os.path.join(config.OUTPUT_DIR, "heatmaps")
        os.makedirs(out_dir, exist_ok=True)
        
        if not positions_log:
            print("[Visualizer] Position log empty. Heatmaps not created.")
            return
            
        df = pd.DataFrame(positions_log)
        
        # Pitch instance
        pitch = Pitch(pitch_type='uefa', pitch_color='#22312b', line_color='#c7d5cc', corner_arcs=True)
        
        for team_id, name in [(0, "team_A"), (1, "team_B")]:
            team_df = df[(df["team_id"] == team_id) & (df["object_id"] != -99)]
            
            if len(team_df) < 5:
                print(f"[Visualizer] Not enough positions logged for {name} to draw heatmap.")
                continue
                
            fig, ax = pitch.draw(figsize=(12, 8))
            
            try:
                # Kernel Density Estimate Plot
                # sns.kdeplot requires x and y. Note: in mplsoccer Pitch, 
                # x is along the width, y is along the height.
                # We can use pitch.kdeplot:
                pitch.kdeplot(
                    team_df["x_pitch"].values, 
                    team_df["y_pitch"].values, 
                    ax=ax, 
                    cmap='hot', 
                    fill=True, 
                    levels=80, 
                    alpha=0.6,
                    zorder=1
                )
                
                ax.set_title(f"Tactical Occupation Density - {name.replace('_', ' ').title()}", 
                             color='white', fontsize=18, pad=15)
                             
                fig.patch.set_facecolor('#22312b')
                
                # Save plot
                filepath = os.path.join(out_dir, f"{name}_heatmap.png")
                fig.savefig(filepath, facecolor='#22312b', edgecolor='none', bbox_inches='tight')
                print(f"[Visualizer] Saved {name} heatmap to {filepath}")
            except Exception as e:
                print(f"[Visualizer] Error creating heatmap for {name}: {e}")
            finally:
                plt.close(fig)