import streamlit as st
import os
import cv2
import time
import tempfile

from ultralytics import YOLO
from huggingface_hub import hf_hub_url, cached_download

@st.cache_resource
def load_model():

    repo_id = 'navoditamathur/Soccer_yolo'
    model_filename = 'soccer_ball.pt'
    # Create a URL for the model file on the Hugging Face Hub
    model_url = hf_hub_url(repo_id, model_filename)

    # Download the model file from the Hub and cache it locally
    cached_model_path = cached_download(model_url)

    # Rename the file to have a .pt extension
    new_cached_model_path = f"{cached_model_path}.pt"
    os.rename(cached_model_path, new_cached_model_path)

    print(f"Downloaded model to {new_cached_model_path}")

    # Load the model using YOLO from the cached model file
    return YOLO(new_cached_model_path)

def process_video(video_path, output_path):
  cap = cv2.VideoCapture(video_path)
  width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  fps = int(cap.get(cv2.CAP_PROP_FPS))
  total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

  fourcc = cv2.VideoWriter_fourcc(*'mp4v')
  out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

  progress_text = "Please wait..."
  progress_bar = st.progress(0)
  progress_bar.text(progress_text)

  status_text = st.empty()
  time_text = st.empty()  

  start_time = time.time()

  for i in range(total_frames):
    ret, frame = cap.read()
    if not ret:
      break

    boxes = model(frame)
    annotated_frame = boxes[0].plot()
    out.write(annotated_frame)

    progress = (i + 1) / total_frames
    progress_bar.progress(progress)

    elasped_time = time.time() - start_time
    time_per_frame = elasped_time / (i + 1)
    remaining_time = (total_frames - (i + 1)) * time_per_frame

    status_text.text(f"Processing frame {i + 1} of {total_frames}")
    time_text.text(f"Time remaining: {remaining_time:.2f} seconds")
  cap.release()
  out.release()
  status_text.text("Video processing completed.")
  progress_bar.empty()
  time_text.empty()


model = load_model()
st.title("Soccer Ball Detection App")

# Sidebar for options
st.sidebar.header("Options")
video_option = st.sidebar.radio("Choose video source:", ("Use preset video", "Upload video"))

if video_option == "Upload video":
    uploaded_file = st.sidebar.file_uploader("Choose a video file", type=["mp4", "avi", "mov"])
    if uploaded_file is not None:
        tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
        tfile.write(uploaded_file.read())
        video_path = tfile.name
else:
        preset_videos = {
            "Soccer Video": "preset_videos/soccer.mp4"
        }
        selected_video = st.sidebar.selectbox("Select a preset video", list(preset_videos.keys()))
        video_path = preset_videos[selected_video]

if 'video_path' in locals():
    st.header("Original Video")
    st.video(video_path)

if st.button("Detect"):
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    output_path = temp_file.name
    process_video(video_path, output_path)
    with open(output_path, 'rb') as video_file:
        video_bytes = video_file.read()
        st.header("Detected Video")
        
        # Debugging: Display video size
        st.write(f"Processed video size: {len(video_bytes)} bytes")

        if len(video_bytes) > 0:
            st.video(video_bytes)
            
        # Generate a download button
        btn = st.download_button(
            label="Download Processed Video",
            data=video_bytes,
            file_name="processed_video.mp4",
            mime="video/mp4"
        )