import os
import cv2
import gradio as gr
from raft import Raft
import yt_dlp

def download_youtube_video(youtube_url, output_filename):
    ydl_opts = {
        'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
        'outtmpl': output_filename,
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        ydl.download([youtube_url])

def process_video(youtube_url, start_time, flow_frame_offset):
    model_path = 'raft_small_iter10_240x320.onnx'
    flow_estimator = Raft(model_path)

    output_filename = 'downloaded_video.mp4'
    processed_output = 'processed_video.mp4'

    # Download video
    if os.path.exists(output_filename):
        os.remove(output_filename)
    download_youtube_video(youtube_url, output_filename)

    cap = cv2.VideoCapture(output_filename)
    if not cap.isOpened():
        return "Error: Could not open video."

    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(processed_output, fourcc, fps, (frame_width, frame_height))

    cap.set(cv2.CAP_PROP_POS_FRAMES, start_time * fps)

    frame_list = []
    frame_num = 0

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

        frame_list.append(prev_frame)
        frame_num += 1

        if frame_num <= flow_frame_offset:
            continue

        flow_map = flow_estimator(frame_list[0], frame_list[-1])
        flow_img = flow_estimator.draw_flow()

        alpha = 0.5
        combined_img = cv2.addWeighted(frame_list[0], alpha, flow_img, (1 - alpha), 0)

        if combined_img is None:
            break

        out.write(combined_img)
        frame_list.pop(0)

    cap.release()
    out.release()

    return processed_output




examples = [
    ["https://www.youtube.com/watch?v=is38pqgbj6A", 5, 50, "output_1.mp4"],
    ["https://www.youtube.com/watch?v=AdbrfoxiAtk", 0, 60, "output_2.mp4"],
    ["https://www.youtube.com/watch?v=vWGg0iPmI8k", 13, 70, "output_3.mp4"],
]

with gr.Blocks() as app:
    gr.Markdown("### Optical Flow Video Processing\n"
                "Enter a YouTube URL, set the start time and flow frame offset, "
                "then click 'Process Video' to see the optical flow processing.")

    with gr.Row():
        with gr.Column():
            youtube_url = gr.Textbox(label="YouTube URL", placeholder="Enter YouTube Video URL Here")
            start_time = gr.Slider(minimum=0, maximum=60, label="Start Time (seconds)", step=1)
            flow_frame_offset = gr.Slider(minimum=1, maximum=100, label="Flow Frame Offset", step=1)
            submit_button = gr.Button("Process Video")

        with gr.Column():
            output_video = gr.Video(label="Processed Video")

    submit_button.click(
        fn=process_video,
        inputs=[youtube_url, start_time, flow_frame_offset],
        outputs=output_video
    )

    gr.Examples(examples=examples,
                 inputs=[youtube_url, start_time, flow_frame_offset],
                 fn=process_video,
                 outputs=output_video,
                 cache_examples=False)

app.launch()