# app.py
import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
import supervision as sv
from scipy.spatial import distance as dist
import tempfile

# Initialize components
model = YOLO('yolov8s.pt')
tracker = sv.ByteTrack()

# Streamlit UI
st.title("⚽️ Player Tracking System")
uploaded_video = st.file_uploader("Upload match video", type=["mp4", "mov"])
calibration_dist = st.number_input("Field width in meters (for speed calibration):", value=68.0)

# Initialize session state
if 'player_data' not in st.session_state:
    st.session_state.player_data = {}

if uploaded_video:
    tfile = tempfile.NamedTemporaryFile(delete=False)
    tfile.write(uploaded_video.read())
    
    cap = cv2.VideoCapture(tfile.name)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    pixels_per_meter = frame_width / calibration_dist if calibration_dist > 0 else 1.0
    
    st_frame = st.empty()
    frame_count = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Detection + tracking
        results = model(frame)[0]
        detections = sv.Detections.from_ultralytics(results)
        detections = tracker.update_with_detections(detections)

        for detection in detections:
            track_id = int(detection[4])
            bbox = detection[0]
            centroid = (int((bbox[0]+bbox[2])/2), int((bbox[1]+bbox[3])/2))
            speed = 0.0  # Initialize speed for all cases

            # Initialize new player
            if track_id not in st.session_state.player_data:
                st.session_state.player_data[track_id] = {
                    'positions': [centroid],
                    'timestamps': [frame_count/fps],
                    'distance': 0.0,
                    'speeds': []
                }
            else:
                # Calculate movement metrics
                prev_pos = st.session_state.player_data[track_id]['positions'][-1]
                time_diff = (frame_count/fps) - st.session_state.player_data[track_id]['timestamps'][-1]
                
                pixel_dist = dist.euclidean(prev_pos, centroid)
                speed_px = pixel_dist / time_diff if time_diff > 0 else 0.0
                speed = speed_px * pixels_per_meter  # Convert to m/s
                
                # Update player record
                st.session_state.player_data[track_id]['positions'].append(centroid)
                st.session_state.player_data[track_id]['timestamps'].append(frame_count/fps)
                st.session_state.player_data[track_id]['distance'] += pixel_dist
                st.session_state.player_data[track_id]['speeds'].append(speed)

            # Annotation with failsafe
            label = f"ID:{track_id} Speed:{speed:.1f}m/s"
            cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), 
                         (int(bbox[2]), int(bbox[3])), (0,255,0), 2)
            cv2.putText(frame, label, (int(bbox[0]), int(bbox[1]-10)),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1)

        st_frame.image(frame, channels="BGR")
        frame_count += 1

    cap.release()
    
    # Display final analytics
    st.subheader("Player Statistics")
    for player_id, data in st.session_state.player_data.items():
        avg_speed = np.mean(data['speeds']) if data['speeds'] else 0
        st.write(f"""
        **Player {player_id}**  
        - Total distance: {data['distance'] * pixels_per_meter:.2f}m  
        - Avg speed: {avg_speed:.1f}m/s  
        - Max speed: {np.max(data['speeds']):.1f}m/s  
        - Active duration: {data['timestamps'][-1]-data['timestamps'][0]:.1f}s
        """)