import gradio as gr
import numpy as np
import cv2

# Function to create a 5x2 video grid
def create_video_grid(video_files, layout, delay):
    # Initialize the video capture objects
    video_captures = [cv2.VideoCapture(video) for video in video_files]
    
    # Get the frame dimensions of the first video
    ret, frame = video_captures[0].read()
    height, width, _ = frame.shape
    
    # Define the output video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter('output.mp4', fourcc, 7, (width * 5, height * 2))
    
    # Define the delay in frames
    delay_frames = int(7 * delay)
    
    # Initialize the frame counters
    frame_counters = [0] * len(video_captures)
    
    while True:
        frames = []
        for i, cap in enumerate(video_captures):
            ret, frame = cap.read()
            if not ret:
                break
            frames.append(frame)
            frame_counters[i] += 1
        
        if len(frames) < 10:
            break
        
        # Apply the layout
        if layout == "Left and Right":
            row1 = np.hstack(frames[:5])
            row2 = np.hstack(frames[5:])
            grid = np.vstack([row1, row2])
        elif layout == "Top and Bottom":
            col1 = np.vstack(frames[:5])
            col2 = np.vstack(frames[5:])
            grid = np.hstack([col1, col2])
        
        # Write the frame to the output video
        out.write(grid)
        
        # Apply the delay for the second row
        if frame_counters[5] < delay_frames:
            for i in range(5, 10):
                frame_counters[i] = 0
                video_captures[i].set(cv2.CAP_PROP_POS_FRAMES, 0)
    
    # Release the video captures and the output video writer
    for cap in video_captures:
        cap.release()
    out.release()
    
    return 'output.mp4'

# Define the Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            video_files = gr.File(label="Upload Video Files", file_count="multiple")
            layout = gr.Radio(choices=["Left and Right", "Top and Bottom"], label="Layout")
            delay = gr.Slider(0, 1, step=0.1, label="Delay (seconds)")
        with gr.Column():
            output_video = gr.Video(label="Output Video", format="mp4")
    
    # Define the event listener
    video_files.change(create_video_grid, [video_files, layout, delay], output_video)

# Launch the Gradio app
demo.launch(show_error=True)