import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from PIL import Image
import io
import os
import cv2
from math import tau
import gradio as gr
from concurrent.futures import ThreadPoolExecutor
import tempfile

def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_kernel_size, desired_range, num_points, theta_points):
    # Convert PIL to OpenCV image
    img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
    img = cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_AREA)
    imgray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(imgray, (blur_kernel_size, blur_kernel_size), 0)
    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # find the contour with the largest area
    largest_contour_idx = np.argmax([cv2.contourArea(c) for c in contours])
    largest_contour = contours[largest_contour_idx]
    
    # def combine_all_contours(contours):
    #     combined_contour = np.array([], dtype=np.int32).reshape(0, 1, 2)
    #     for contour in contours:
    #         combined_contour = np.vstack((combined_contour, contour))
    #     return combined_contour

    # largest_contour = combine_all_contours(contours)
    verts = [tuple(coord) for coord in largest_contour.squeeze()]
    xs, ys = np.asarray(list(zip(*verts)))
    x_range, y_range = np.max(xs) - np.min(xs), np.max(ys) - np.min(ys)
    scale_x, scale_y = desired_range / x_range, desired_range / y_range
    xs = (xs - np.mean(xs)) * scale_x
    ys = (-ys + np.mean(ys)) * scale_y

    t_list = np.linspace(0, tau, len(xs))
    t_values = np.linspace(0, tau, num_points)
    f_precomputed = np.interp(t_values, t_list, xs + 1j * ys)

    def compute_cn(f_exp, n, t_values):
        coef = np.trapz(f_exp * np.exp(-n * t_values * 1j), t_values) / tau
        return coef

    N = coefficients
    indices = [0] + [j for i in range(1, N + 1) for j in (i, -i)]

    with ThreadPoolExecutor(max_workers=8) as executor:
        coefs = list(executor.map(lambda n: (compute_cn(f_precomputed, n, t_values), n), indices))

    fig, ax = plt.subplots()
    circles = [ax.plot([], [], 'b-')[0] for _ in range(-N, N + 1)]
    circle_lines = [ax.plot([], [], 'g-')[0] for _ in range(-N, N + 1)]
    drawing, = ax.plot([], [], 'r-', linewidth=2)

    ax.set_xlim(-desired_range, desired_range)
    ax.set_ylim(-desired_range, desired_range)
    ax.set_axis_off()
    ax.set_aspect('equal')
    fig.set_size_inches(15, 15)

    draw_x, draw_y = [], []
    theta = np.linspace(0, tau, theta_points)
    coefs_static = [(np.linalg.norm(c), fr) for c, fr in coefs]
    last_image = None 
    
    # Initialize the background
    fig.canvas.draw()
    background = fig.canvas.copy_from_bbox(ax.bbox)
    
    def animate(i, coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta):
        # Restore the background to erase old frames
        fig.canvas.restore_region(background)
        
        center = (0, 0)
        for idx, (r, fr) in enumerate(coefs_static):
            c_dynamic = coefs[idx][0] * np.exp(1j * (fr * tau * time[i]))
            x, y = center[0] + r * np.cos(theta), center[1] + r * np.sin(theta)
            circle_lines[idx].set_data([center[0], center[0] + np.real(c_dynamic)], [center[1], center[1] + np.imag(c_dynamic)])
            circles[idx].set_data(x, y)
            center = (center[0] + np.real(c_dynamic), center[1] + np.imag(c_dynamic))
    
        draw_x.append(center[0])
        draw_y.append(center[1])
        drawing.set_data(draw_x[:i+1], draw_y[:i+1])
    
        # Draw only the updated elements
        for circle in circles:
            ax.draw_artist(circle)
        for line in circle_lines:
            ax.draw_artist(line)
        ax.draw_artist(drawing)
    
        # Blit only the updated area
        fig.canvas.blit(ax.bbox)
    
        # Capture the current canvas state as a PIL Image
        canvas = fig.canvas
        w, h = canvas.get_width_height()
        buf = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8)
        image = Image.fromarray(buf.reshape(h, w, 4), 'RGBA').convert('RGB')
        
        return (image, None)

    # Generate and yield images for each frame
    time = np.linspace(0, 1, num=frames)
    for frame in range(frames):
        yield from animate(frame, coefs, )
        yield from animate(frame, coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta)

    # Generate final animation
    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file:
        anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, np.linspace(0, 1, num=frames)))
        anim.save(temp_file.name, fps=15)

    yield (last_image, temp_file.name)

# Gradio interface setup
interface = gr.Interface(
    fn=fourier_transform_drawing,
    inputs=[
        gr.Image(label="Input Image", sources=['upload'], type="pil"),
        gr.Slider(minimum=5, maximum=500, value=100, label="Number of Frames"),
        gr.Slider(minimum=1, maximum=500, value=50, label="Number of Coefficients"),
        gr.Number(value=224, label="Image Size (px)", precision=0),
        gr.Slider(minimum=3, maximum=11, step=2, value=5, label="Blur Kernel Size (odd number)"),
        gr.Number(value=400, label="Desired Range for Scaling", precision=0),
        gr.Number(value=1000, label="Number of Points for Integration", precision=0),
        gr.Slider(minimum=50, maximum=500, value=80, label="Theta Points for Animation")
    ],
    outputs=["image", gr.Video()],
    title="Fourier Transform Drawing",
    description="Upload an image and generate a Fourier Transform drawing animation.",
)

if __name__ == "__main__":
    # define queue - required for generators
    interface.queue()
    interface.launch()