import streamlit as st
import cv2
import tempfile
import numpy as np
from ultralytics import YOLO, solutions

# Load the YOLOv8 model
model = YOLO("yolov8n-pose.pt")

# Streamlit App
st.title("Workout Monitoring App")
st.write("Upload a video to monitor your ab workout.")

uploaded_file = st.file_uploader("Choose a video file", type=["mp4", "mov", "avi"])

if uploaded_file is not None:
    # Save the uploaded video to a temporary file
    tfile = tempfile.NamedTemporaryFile(delete=False)
    tfile.write(uploaded_file.read())
    tfile.close()

    # Load the video with OpenCV
    cap = cv2.VideoCapture(tfile.name)
    assert cap.isOpened(), "Error reading video file"
    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # Initialize AIGym object
    gym_object = solutions.AIGym(
        line_thickness=2,
        view_img=False,  # Set to False since we are using Streamlit to display
        pose_type="abworkout",  # Use 'abworkout' as the pose type
        kpts_to_check=[6, 8, 10],
    )

    # List to store processed frames
    processed_frames = []

    # Streamlit progress bar
    progress_bar = st.progress(0)
    frame_count = 0

    # Process the video frame by frame
    st.write("Analyzing video... Please wait.")
    while cap.isOpened():
        success, im0 = cap.read()
        if not success:
            break

        results = model.track(im0, verbose=False)  # Tracking recommended
        im0 = gym_object.start_counting(im0, results)

        # Resize the frame to 320x320
        im0_resized = cv2.resize(im0, (1024, 1024))

        # Append the processed frame to the list
        processed_frames.append(im0_resized)

        # Update progress bar
        frame_count += 1
        progress_bar.progress(frame_count / total_frames)

    cap.release()
    cv2.destroyAllWindows()

    # Create a temporary file to save the processed video
    output_video_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    output_video_path = output_video_file.name

    # Write the processed frames to a video file
    video_writer = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (1024, 1024))
    for frame in processed_frames:
        video_writer.write(frame)
    video_writer.release()

    # Display the processed video in Streamlit
    st.write("Analysis complete. Displaying processed video:")
    st.video(output_video_path)

    # Provide a download link for the processed video
    st.write("Download the processed video:")
    with open(output_video_path, "rb") as video_file:
        st.download_button("Download", video_file, "workouts.mp4")