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_video)
    # Read Video
    video_frames = read_video(input_video)  # use the uploaded file
    
    # Initialize Tracker
    tracker = Tracker('models/best.pt')

    if input_video.contains("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.contains("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 tool processes football game videos to track player movements, team assignments, ball possession, speed, distance, and camera movement. 

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

    Original Reference: https://www.youtube.com/watch?v=neBZ6huolkg"""

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)