import gradio as gr
import numpy as np
import cv2

def create_dot_effect(image, dot_size=10, spacing=2, invert=False):
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        gray = image
    
    gray = cv2.adaptiveThreshold(
        gray,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        25,
        5
    )
    
    height, width = gray.shape
    canvas = np.zeros_like(gray) if not invert else np.full_like(gray, 255)
    
    y_dots = range(0, height, dot_size + spacing)
    x_dots = range(0, width, dot_size + spacing)
    
    dot_color = 255 if not invert else 0
    for y in y_dots:
        for x in x_dots:
            region = gray[y:min(y+dot_size, height), x:min(x+dot_size, width)]
            if region.size > 0:
                brightness = np.mean(region)
                relative_brightness = brightness / 255.0
                if invert:
                    relative_brightness = 1 - relative_brightness
                radius = int((dot_size/2) * relative_brightness)
                if radius > 0:
                    cv2.circle(canvas, 
                             (x + dot_size//2, y + dot_size//2), 
                             radius, 
                             (dot_color), 
                             -1)
    return canvas

def process_video(video_path, dot_size=10, spacing=2, invert=False, fps=30):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None

    orig_fps = int(cap.get(cv2.CAP_PROP_FPS))
    orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # Ensure output FPS does not exceed input FPS
    output_fps = min(fps, orig_fps)
    
    max_height = 720
    if orig_height > max_height:
        scale = max_height / orig_height
        frame_width = int(orig_width * scale)
        frame_height = max_height
    else:
        frame_width = orig_width
        frame_height = orig_height
    
    # Ensure even dimensions for video codec compatibility
    frame_width = frame_width // 2 * 2
    frame_height = frame_height // 2 * 2
    
    output_path = "output.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, output_fps, (frame_width, frame_height), True)

    try:
        frame_counter = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            
            # Skip or duplicate frames based on FPS ratio
            if output_fps < orig_fps:
                # Skip frames to match the output FPS
                if frame_counter % (orig_fps // output_fps) != 0:
                    frame_counter += 1
                    continue
            elif output_fps > orig_fps:
                # Duplicate frames to match the output FPS
                for _ in range(output_fps // orig_fps):
                    if frame.shape[0] != frame_height:
                        frame = cv2.resize(frame, (frame_width, frame_height))
                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    dotted_frame = create_dot_effect(frame_rgb, dot_size, spacing, invert)
                    dotted_frame_bgr = cv2.cvtColor(dotted_frame, cv2.COLOR_GRAY2BGR)
                    out.write(dotted_frame_bgr)
            
            if frame.shape[0] != frame_height:
                frame = cv2.resize(frame, (frame_width, frame_height))
            
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            dotted_frame = create_dot_effect(frame_rgb, dot_size, spacing, invert)
            dotted_frame_bgr = cv2.cvtColor(dotted_frame, cv2.COLOR_GRAY2BGR)
            out.write(dotted_frame_bgr)
            
            frame_counter += 1

    finally:
        cap.release()
        out.release()
    
    return output_path

def get_video_fps(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return 30  # Default FPS if video cannot be opened
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    cap.release()
    return fps

with gr.Blocks(title="Dot Effect Generator") as iface:
    gr.Markdown("# 🎨 Dot Effect Generator")
    gr.Markdown("Transform media into stylized dot patterns")
    
    with gr.Tab("🖼️ Image"):
        with gr.Row():
            img_input = gr.Image(label="Input")
            img_output = gr.Image(label="Output")
        with gr.Row():
            img_dot_size = gr.Slider(2, 20, 10, step=1, label="Dot Size")
            img_spacing = gr.Slider(0, 10, 2, step=1, label="Spacing")
        img_btn = gr.Button("Generate", variant="primary")
        img_btn.click(create_dot_effect, [img_input, img_dot_size, img_spacing], img_output)
    
    with gr.Tab("🎥 Video"):
        with gr.Row():
            vid_input = gr.Video(label="Input", format="mp4")
            vid_output = gr.Video(label="Output", format="mp4")
        with gr.Row():
            vid_dot_size = gr.Slider(2, 20, 10, step=1, label="Dot Size")
            vid_spacing = gr.Slider(0, 10, 2, step=1, label="Spacing")
            vid_invert = gr.Checkbox(label="Invert Colors")
        with gr.Row():
            vid_fps = gr.Slider(1, 60, 30, step=1, label="Output FPS")
        vid_btn = gr.Button("Process Video", variant="primary")
        
        # Update FPS slider max value based on input video
        def update_fps_slider(video_path):
            if video_path is None:
                return gr.update(maximum=60, value=30)  # Default if no video
            fps = get_video_fps(video_path)
            return gr.update(maximum=fps, value=min(fps, 30))
        
        vid_input.change(update_fps_slider, vid_input, vid_fps)
        
        vid_btn.click(process_video, [vid_input, vid_dot_size, vid_spacing, vid_invert, vid_fps], vid_output)

if __name__ == "__main__":
    iface.launch(server_port=7860, show_error=True)