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

def process_image(input_image, img_size, blur_kernel_size, desired_range):
    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)

    largest_contour_idx = np.argmax([cv2.contourArea(c) for c in contours])
    largest_contour = contours[largest_contour_idx]
    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

    return xs, 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

def calculate_fourier_coefficients(xs, ys, num_points, coefficients):
    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)

    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), indices))
    
    return coefs

def setup_animation_env(img_size, desired_range, coefficients):
    fig, ax = plt.subplots()
    circles = [ax.plot([], [], 'b-')[0] for _ in range(-coefficients, coefficients + 1)]
    circle_lines = [ax.plot([], [], 'g-')[0] for _ in range(-coefficients, coefficients + 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)
    fig.canvas.draw()
    background = fig.canvas.copy_from_bbox(ax.bbox)

    return fig, ax, background, circles, circle_lines, drawing

def animate(frame, coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta):
    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[frame]))
        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, draw_y)

    for circle in circles:
        ax.draw_artist(circle)
    for line in circle_lines:
        ax.draw_artist(line)
    ax.draw_artist(drawing)

    fig.canvas.blit(ax.bbox)

def generate_animation(frames, coefs, img_size, desired_range, theta_points, coefficients):
    fig, ax, background, circles, circle_lines, drawing = setup_animation_env(img_size, desired_range, coefficients)
    coefs_static = [(np.linalg.norm(c), fr) for c, fr in coefs]
    time = np.linspace(0, 1, num=frames)
    theta = np.linspace(0, tau, theta_points)
    draw_x, draw_y = [], []

    anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta))

    return anim

def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_kernel_size, desired_range, num_points, theta_points):
    xs, ys = process_image(input_image, img_size, blur_kernel_size, desired_range)
    coefs = calculate_fourier_coefficients(xs, ys, num_points, coefficients)
    anim = generate_animation(frames, coefs, img_size, desired_range, theta_points, coefficients)
    
    # Saving the animation
    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file:
        anim.save(temp_file.name, fps=15)
        return Image.fromarray(np.zeros((img_size, img_size), dtype=np.uint8)), temp_file.name

def setup_gradio_interface():
    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."
    )
    return interface

if __name__ == "__main__":
    interface = setup_gradio_interface()
    interface.queue()
    interface.launch()