modules/analytics.py

import os
import json
import sqlite3
import pandas as pd
import numpy as np
import football_analytics.config as config

class FootballAnalytics:
    def __init__(self, fps=30):
        self.fps = fps
        
        # Data storage
        self.positions_log = []     # Raw tracking data: frame, time, track_id, team_id, x_img, y_img, x_pitch, y_pitch
        self.possession_log = []    # Frame-by-frame possession: frame, player_id, team_id
        self.events_log = []        # Event data: type (pass, interception, etc.), frame, player_id, team_id, extra info
        self.passes_log = []        # Pass specific data: frame, passer_id, passer_team, receiver_id, receiver_team, success, distance
        
        # Player metrics tracking
        self.player_stats = {}      # track_id -> {team_id, distance, frames_played, positions_list (pitch), successful_passes, total_passes}
        self.last_positions = {}    # track_id -> (x_pitch, y_pitch) of previous frame
        
        # Possession state machine
        self.current_possession_player = None
        self.current_possession_team = None
        self.candidate_player = None
        self.candidate_frames = 0
        self.ball_out_of_contact_frames = 0
        
        # Last known ball position for interpolation/memory
        self.last_ball_pos = None
        self.last_ball_frame = -1
        
        # Zone tracking grid (3x6)
        self.grid_rows = config.GRID_ROWS
        self.grid_cols = config.GRID_COLS
        self.zone_counts = {0: np.zeros(self.grid_rows * self.grid_cols), 
                            1: np.zeros(self.grid_rows * self.grid_cols)}

    def _get_zone_id(self, x, y):
        """
        Maps a 2D pitch coordinate (x, y) to a zone ID (0 to 17 for a 3x6 grid).
        """
        # Clip coordinates to pitch dimensions
        x = max(0.0, min(x, config.PITCH_WIDTH))
        y = max(0.0, min(y, config.PITCH_HEIGHT))
        
        col = int(x / (config.PITCH_WIDTH / self.grid_cols))
        row = int(y / (config.PITCH_HEIGHT / self.grid_rows))
        
        col = min(col, self.grid_cols - 1)
        row = min(row, self.grid_rows - 1)
        
        return row * self.grid_cols + col

    def update_frame(self, frame_num, detections, team_assignments, ball_detection=None):
        """
        Updates the analytics state for a new frame.
        
        Args:
            frame_num (int): Current frame index
            detections (sv.Detections): Tracked player detections
            team_assignments (dict): Mapping tracker_id -> team_id (0: Team A, 1: Team B, 2: Referee/Other)
            ball_detection (tuple): (x_pitch, y_pitch, x_img, y_img) or None
        """
        timestamp = frame_num / self.fps
        
        # 1. Parse player detections and update stats
        players_in_frame = {}
        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 or track_id == -1:  # Only process tracked persons (players)
                continue
                
            team_id = team_assignments.get(track_id, 0)
            if team_id == 2:  # Referee/Other, skip from tactical statistics
                continue
                
            # Bottom-center of bbox as feet position
            x_img = (bbox[0] + bbox[2]) / 2.0
            y_img = bbox[3]
            
            # Retrieve projected coordinates
            # Since these have already been projected by pipeline, let's assume we pass pitch coords in extra properties
            # Or we can compute it on the fly. We'll pass them from pipeline.
            # We will store them in positions_log.
            
        # 2. Update distance covered
        # We will handle coordinate conversion in pipeline.py and log positions via log_positions()
        pass

    def log_positions(self, frame_num, player_data, ball_data):
        """
        Logs player and ball positions, and updates cumulative metrics like distance.
        
        Args:
            frame_num (int): Current frame index
            player_data (list of dict): [{id, team, x_img, y_img, x_pitch, y_pitch}]
            ball_data (dict or None): {x_img, y_img, x_pitch, y_pitch} or None
        """
        timestamp = frame_num / self.fps
        
        # Record ball position
        if ball_data is not None:
            self.last_ball_pos = (ball_data["x_pitch"], ball_data["y_pitch"])
            self.last_ball_frame = frame_num
            self.positions_log.append({
                "frame": frame_num,
                "timestamp": timestamp,
                "object_id": -99,  # Conventional ID for the ball
                "team_id": -99,
                "x_pixel": ball_data["x_img"],
                "y_pixel": ball_data["y_img"],
                "x_pitch": ball_data["x_pitch"],
                "y_pitch": ball_data["y_pitch"]
            })
            self.ball_out_of_contact_frames = 0
        else:
            self.ball_out_of_contact_frames += 1
            if self.ball_out_of_contact_frames > 30:  # Ball lost for more than 1 sec
                self.last_ball_pos = None

        # Record players
        for p in player_data:
            pid = p["id"]
            team = p["team"]
            x_pitch = p["x_pitch"]
            y_pitch = p["y_pitch"]
            
            # Log raw position
            self.positions_log.append({
                "frame": frame_num,
                "timestamp": timestamp,
                "object_id": pid,
                "team_id": team,
                "x_pixel": p["x_img"],
                "y_pixel": p["y_img"],
                "x_pitch": x_pitch,
                "y_pitch": y_pitch
            })
            
            # Initialize stats if first time seeing player
            if pid not in self.player_stats:
                self.player_stats[pid] = {
                    "team_id": team,
                    "distance": 0.0,
                    "frames_played": 0,
                    "positions_list": [],
                    "successful_passes": 0,
                    "total_passes": 0
                }
                
            stats = self.player_stats[pid]
            stats["frames_played"] += 1
            stats["positions_list"].append((x_pitch, y_pitch))
            
            # Calculate distance traveled
            if pid in self.last_positions:
                prev_x, prev_y = self.last_positions[pid]
                dist = np.sqrt((x_pitch - prev_x)**2 + (y_pitch - prev_y)**2)
                
                # Check for unrealistic speed jump (e.g. tracking ID swap)
                # Max running speed is ~10 m/s. Frame time is 1/30s. Max distance per frame is 0.33m.
                # Let's cap frame distance to 1.5m to avoid extreme track jumps affecting analytics.
                if dist < 1.5:
                    stats["distance"] += dist
                    
            self.last_positions[pid] = (x_pitch, y_pitch)
            
            # Record zone occupation
            zone_id = self._get_zone_id(x_pitch, y_pitch)
            self.zone_counts[team][zone_id] += 1

        # Calculate Possession and Passes for this frame
        self._update_possession(frame_num, player_data)
        
        # Calculate Marking
        self._calculate_marking(frame_num, player_data)

    def _update_possession(self, frame_num, player_data):
        """
        Calculates ball possession and logs pass/interception events.
        """
        if self.last_ball_pos is None or len(player_data) == 0:
            return

        bx, by = self.last_ball_pos
        
        # Find player closest to the ball
        closest_player = None
        min_dist = float('inf')
        
        for p in player_data:
            px, py = p["x_pitch"], p["y_pitch"]
            dist = np.sqrt((px - bx)**2 + (py - by)**2)
            if dist < min_dist:
                min_dist = dist
                closest_player = p

        # Check if closest player is within possession range
        if min_dist <= config.POSSESSION_DISTANCE_THRESHOLD:
            pid = closest_player["id"]
            team = closest_player["team"]
            
            # Possession Stabilization Logic (anti-noise)
            if pid == self.candidate_player:
                self.candidate_frames += 1
            else:
                self.candidate_player = pid
                self.candidate_frames = 1
                
            # DEBUG print (can be uncommented for profiling)
            # print(f"Frame {frame_num}: Closest player {pid} (team {team}) dist {min_dist:.2f}m. Candidate frames: {self.candidate_frames}")
                
            if self.candidate_frames >= config.POSSESSION_STABILIZATION_FRAMES:
                # We have stabilized possession!
                if self.current_possession_player != pid:
                    # Log possession change event
                    prev_player = self.current_possession_player
                    prev_team = self.current_possession_team
                    
                    # print(f"Frame {frame_num}: Possession changing from {prev_player} (team {prev_team}) to {pid} (team {team})!")
                    
                    self.current_possession_player = pid
                    self.current_possession_team = team
                    
                    self.possession_log.append({
                        "frame": frame_num,
                        "timestamp": frame_num / self.fps,
                        "player_id": pid,
                        "team_id": team
                    })
                    
                    # Detect if a pass or interception occurred
                    if prev_player is not None and prev_player != pid:
                        # Ball must have traveled between them
                        # Check pass success
                        success = (prev_team == team)
                        dist_pass = np.sqrt((bx - self.last_positions.get(prev_player, (bx, by))[0])**2 + 
                                            (by - self.last_positions.get(prev_player, (bx, by))[1])**2)
                        
                        # Increment stats
                        if prev_player in self.player_stats:
                            self.player_stats[prev_player]["total_passes"] += 1
                            if success:
                                self.player_stats[prev_player]["successful_passes"] += 1
                                
                        event_type = "pass" if success else "interception"
                        
                        event_desc = {
                            "event_id": len(self.events_log) + 1,
                            "frame": frame_num,
                            "timestamp": frame_num / self.fps,
                            "type": event_type,
                            "player_id": prev_player,
                            "team_id": prev_team,
                            "x": bx,
                            "y": by,
                            "receiver_id": pid,
                            "receiver_team_id": team,
                            "result": "success" if success else "failed"
                        }
                        self.events_log.append(event_desc)
                        
                        self.passes_log.append({
                            "frame": frame_num,
                            "timestamp": frame_num / self.fps,
                            "passer_id": prev_player,
                            "passer_team": prev_team,
                            "receiver_id": pid,
                            "receiver_team": team,
                            "success": int(success),
                            "distance": float(dist_pass)
                        })
        else:
            # Ball is loose
            # We don't clear possession immediately unless the ball has been far from everyone for a while
            self.candidate_player = None
            self.candidate_frames = 0

    def _calculate_marking(self, frame_num, player_data):
        """
        For each attacking player (team in possession), calculates the closest defender and distance.
        """
        if self.current_possession_team is None:
            return
            
        attacking_team = self.current_possession_team
        defending_team = 1 - attacking_team
        
        attackers = [p for p in player_data if p["team"] == attacking_team]
        defenders = [p for p in player_data if p["team"] == defending_team]
        
        if not attackers or not defenders:
            return
            
        # We don't log marking per-frame in a separate CSV (too large),
        # but we can track it to find average marking tightess or log it inside SQLite if needed.
        # Let's save average marking distances in summary stats.
        pass

    def get_summary(self):
        """
        Computes final match stats summary.
        """
        # Calculate possession percentages
        total_possession_frames = len(self.possession_log)
        team0_pos_pct = 50.0
        team1_pos_pct = 50.0
        
        if total_possession_frames > 0:
            team0_frames = sum(1 for p in self.possession_log if p["team_id"] == 0)
            team0_pos_pct = (team0_frames / total_possession_frames) * 100.0
            team1_pos_pct = 100.0 - team0_pos_pct

        # Team stats
        team_stats = {
            0: {"distance": 0.0, "passes_attempted": 0, "passes_completed": 0},
            1: {"distance": 0.0, "passes_attempted": 0, "passes_completed": 0}
        }
        
        player_summaries = []
        for pid, stats in self.player_stats.items():
            team = stats["team_id"]
            dist = stats["distance"]
            
            # Update team totals
            team_stats[team]["distance"] += dist
            team_stats[team]["passes_attempted"] += stats["total_passes"]
            team_stats[team]["passes_completed"] += stats["successful_passes"]
            
            # Position average
            avg_x = 0.0
            avg_y = 0.0
            if stats["positions_list"]:
                positions = np.array(stats["positions_list"])
                avg_x = float(positions[:, 0].mean())
                avg_y = float(positions[:, 1].mean())
                
            player_summaries.append({
                "player_id": pid,
                "team_id": team,
                "distance_covered_meters": round(dist, 2),
                "possession_time_seconds": round(sum(1 for p in self.possession_log if p["player_id"] == pid) / self.fps, 2),
                "successful_passes": stats["successful_passes"],
                "total_passes": stats["total_passes"],
                "pass_accuracy": round((stats["successful_passes"] / stats["total_passes"] * 100.0), 1) if stats["total_passes"] > 0 else 0.0,
                "average_x": round(avg_x, 2),
                "average_y": round(avg_y, 2)
            })

        # Heatmap / Zone summaries
        zone_sum0 = self.zone_counts[0].tolist()
        zone_sum1 = self.zone_counts[1].tolist()
        
        # Normalize zones to percentages
        if sum(zone_sum0) > 0:
            zone_sum0 = [round((z / sum(zone_sum0)) * 100, 1) for z in zone_sum0]
        if sum(zone_sum1) > 0:
            zone_sum1 = [round((z / sum(zone_sum1)) * 100, 1) for z in zone_sum1]

        summary = {
            "match_statistics": {
                "possession_percent_team_A": round(team0_pos_pct, 1),
                "possession_percent_team_B": round(team1_pos_pct, 1),
                "total_distance_team_A_meters": round(team_stats[0]["distance"], 1),
                "total_distance_team_B_meters": round(team_stats[1]["distance"], 1),
                "passes_attempted_team_A": team_stats[0]["passes_attempted"],
                "passes_completed_team_A": team_stats[0]["passes_completed"],
                "pass_accuracy_team_A": round((team_stats[0]["passes_completed"] / team_stats[0]["passes_attempted"] * 100.0), 1) if team_stats[0]["passes_attempted"] > 0 else 0.0,
                "passes_attempted_team_B": team_stats[1]["passes_attempted"],
                "passes_completed_team_B": team_stats[1]["passes_completed"],
                "pass_accuracy_team_B": round((team_stats[1]["passes_completed"] / team_stats[1]["passes_attempted"] * 100.0), 1) if team_stats[1]["passes_attempted"] > 0 else 0.0,
            },
            "team_A_zone_occupancy_percent": zone_sum0,
            "team_B_zone_occupancy_percent": zone_sum1,
            "players": player_summaries
        }
        
        return summary

    def export_results(self, output_dir=None):
        """
        Exports the results to CSV, JSON, and SQLite.
        """
        out_dir = output_dir or config.OUTPUT_DIR
        os.makedirs(out_dir, exist_ok=True)
        
        # 1. Export Positions CSV
        df_pos = pd.DataFrame(self.positions_log)
        pos_csv_path = os.path.join(out_dir, "positions.csv")
        df_pos.to_csv(pos_csv_path, index=False)
        print(f"[Analytics] Exported positions to {pos_csv_path}")
        
        # 2. Export Events CSV
        df_events = pd.DataFrame(self.events_log)
        events_csv_path = os.path.join(out_dir, "events.csv")
        df_events.to_csv(events_csv_path, index=False)
        print(f"[Analytics] Exported events to {events_csv_path}")
        
        # 3. Export Passes CSV
        df_passes = pd.DataFrame(self.passes_log)
        passes_csv_path = os.path.join(out_dir, "passes.csv")
        df_passes.to_csv(passes_csv_path, index=False)
        print(f"[Analytics] Exported passes to {passes_csv_path}")
        
        # 4. Export Players CSV
        summary = self.get_summary()
        df_players = pd.DataFrame(summary["players"])
        players_csv_path = os.path.join(out_dir, "players.csv")
        df_players.to_csv(players_csv_path, index=False)
        print(f"[Analytics] Exported player stats to {players_csv_path}")
        
        # 5. Export Summary JSON
        json_path = os.path.join(out_dir, "summary.json")
        with open(json_path, 'w') as f:
            json.dump(summary, f, indent=4)
        print(f"[Analytics] Exported summary JSON to {json_path}")
        
        # 6. Export to SQLite
        self.export_to_sqlite(out_dir)

    def export_to_sqlite(self, out_dir):
        """
        Saves all tables into a SQLite database.
        """
        db_path = os.path.join(out_dir, "football_analytics.sqlite")
        conn = sqlite3.connect(db_path)
        
        try:
            # Drop tables if exist
            cursor = conn.cursor()
            cursor.execute("DROP TABLE IF EXISTS positions")
            cursor.execute("DROP TABLE IF EXISTS events")
            cursor.execute("DROP TABLE IF EXISTS passes")
            cursor.execute("DROP TABLE IF EXISTS players")
            cursor.execute("DROP TABLE IF EXISTS game_summary")
            conn.commit()
            
            # Write pandas dataframes (safely checking if logs are empty to avoid empty schema SQL errors)
            if self.positions_log:
                pd.DataFrame(self.positions_log).to_sql("positions", conn, index=False, if_exists="replace")
            else:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS positions (
                        frame INTEGER, timestamp REAL, object_id INTEGER, team_id INTEGER,
                        x_pixel REAL, y_pixel REAL, x_pitch REAL, y_pitch REAL
                    )
                """)

            if self.events_log:
                pd.DataFrame(self.events_log).to_sql("events", conn, index=False, if_exists="replace")
            else:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS events (
                        event_id INTEGER, frame INTEGER, timestamp REAL, type TEXT,
                        player_id INTEGER, team_id INTEGER, x REAL, y REAL,
                        receiver_id INTEGER, receiver_team_id INTEGER, result TEXT
                    )
                """)

            if self.passes_log:
                pd.DataFrame(self.passes_log).to_sql("passes", conn, index=False, if_exists="replace")
            else:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS passes (
                        frame INTEGER, timestamp REAL, passer_id INTEGER, passer_team INTEGER,
                        receiver_id INTEGER, receiver_team INTEGER, success INTEGER, distance REAL
                    )
                """)
            
            summary = self.get_summary()
            if summary["players"]:
                pd.DataFrame(summary["players"]).to_sql("players", conn, index=False, if_exists="replace")
            else:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS players (
                        player_id INTEGER, team_id INTEGER, distance_covered_meters REAL,
                        possession_time_seconds REAL, successful_passes INTEGER, total_passes INTEGER,
                        pass_accuracy REAL, average_x REAL, average_y REAL
                    )
                """)
            
            # Match stats table
            match_stats = summary["match_statistics"]
            df_match = pd.DataFrame([match_stats])
            df_match.to_sql("game_summary", conn, index=False, if_exists="replace")
            
            conn.commit()
            print(f"[Analytics] Exported tables to SQLite database at {db_path}")
        except Exception as e:
            print(f"[Analytics] Error exporting to SQLite: {e}")
        finally:
            conn.close()