import gradio as gr

import os
import numpy as np

from utils import read_video, save_video
from trackers import Tracker
from team_assigner import TeamAssigner
from player_ball_assigner import PlayerBallAssigner
from camera_movement_estimator import CameraMovementEstimator
from view_transformer import ViewTransformer
from speed_and_distance_estimator import SpeedAndDistance_Estimator

def process_video(input_video, player_stats=True, ball_stats=True):
    print("input: " + input_video)
    # Read Video
    video_frames = read_video(input_video)  # use the uploaded file
    
    # Initialize Tracker
    tracker = Tracker('models/best.pt')

    if input_video.endswith("121364_0_small.mp4"):
        print("loading cached tracks")
        tracks = tracker.get_object_tracks(video_frames,
                                       read_from_stub=True,
                                       stub_path='stubs/track_stub_121364_0_small.pkl')
    else: 
        tracks = tracker.get_object_tracks(video_frames)

    # Interpolate Ball Positions
    tracks["ball"] = tracker.interpolate_ball_positions(tracks["ball"])

    # Get object positions
    tracker.add_position_to_tracks(tracks)

    # Camera movement estimator
    camera_movement_estimator = CameraMovementEstimator(video_frames[0])
    if input_video.endswith("121364_0_small.mp4"):
        print("loading cached camera movements")
        camera_movement_per_frame = camera_movement_estimator.get_camera_movement(video_frames,
                                                                                read_from_stub=True,
                                                                                stub_path='stubs/camera_movement_stub_121364_0_small.pkl')
    else: 
        camera_movement_per_frame = camera_movement_estimator.get_camera_movement(video_frames)
    camera_movement_estimator.add_adjust_positions_to_tracks(tracks, camera_movement_per_frame)

    # View Transformer
    view_transformer = ViewTransformer()
    view_transformer.add_transformed_position_to_tracks(tracks)

    # Speed and distance estimator
    speed_and_distance_estimator = SpeedAndDistance_Estimator()
    exclude_objects=['referees', 'players', 'ball']
    if player_stats:
        exclude_objects.remove('players')
    if ball_stats:
        exclude_objects.remove('ball')
    speed_and_distance_estimator.add_speed_and_distance_to_tracks(tracks, exclude_objects)

    # Assign Player Teams
    team_assigner = TeamAssigner()
    team_assigner.assign_team_color(video_frames[0], tracks['players'][0])

    for frame_num, player_track in enumerate(tracks['players']):
        for player_id, track in player_track.items():
            team = team_assigner.get_player_team(video_frames[frame_num], track['bbox'], player_id)
            tracks['players'][frame_num][player_id]['team'] = team
            tracks['players'][frame_num][player_id]['team_color'] = team_assigner.team_colors[team]

    # Assign Ball Acquisition
    player_assigner = PlayerBallAssigner()
    team_ball_control = []
    for frame_num, player_track in enumerate(tracks['players']):
        ball_bbox = tracks['ball'][frame_num][1]['bbox']
        assigned_player = player_assigner.assign_ball_to_player(player_track, ball_bbox)

        if assigned_player != -1:
            tracks['players'][frame_num][assigned_player]['has_ball'] = True
            team_ball_control.append(tracks['players'][frame_num][assigned_player]['team'])
        else:
            if team_ball_control: # in case first few frames assigned_player == -1
                team_ball_control.append(team_ball_control[-1])
    team_ball_control = np.array(team_ball_control)

    # Draw output
    output_video_frames = tracker.draw_annotations(video_frames, tracks, team_ball_control)
    output_video_frames = camera_movement_estimator.draw_camera_movement(output_video_frames, camera_movement_per_frame)
    speed_and_distance_estimator.draw_speed_and_distance(output_video_frames, tracks)

    # Save output video
    output_path = 'output_videos/output_video.avi'
    save_video(output_video_frames, output_path)

    return output_path

# Gradio Interface
title="Football Match Analytics with YOLO and OpenCV"
description="""
    This demo processes football game videos to detect players and referees, track the ball, assign players to teams using color pixel clustering, and compute ball possession per team. 
    
    It also estimates camera movement with Lucas-Kanade optical flow and applies perspective transformation to calculate the real-time speed and total distance traveled by each player and the ball.

    The YOLO detection model was fine-tuned with this dataset: https://universe.roboflow.com/roboflow-jvuqo/football-players-detection-3zvbc/dataset

    Original Tutorial Reference: https://www.youtube.com/watch?v=neBZ6huolkg
    
    **Note**: this space is running on a CPU, so inferencing new video may take some time. (Avg time during test: 1min processing per 5 second of video)"""

examples = [["input_videos/121364_0_small.mp4", True, True]]

interface = gr.Interface(fn=process_video,
                            inputs=[
                                gr.Video(label="Upload Video (mp4, avi, mov) Max: 30sec"),
                                gr.Checkbox(label="Include Player Stats", value=True),
                                gr.Checkbox(label="Include Ball Stats", value=True)
                                ],
                            outputs=gr.Video(label="Processed Video"),
                            examples=examples,
                            live=False,  # No live update to avoid real-time processing issues
                            title=title,
                            description=description)  # Allow users to download the processed video

interface.launch(debug=True, show_error = True)