import gradio as gr
import cv2
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import quad_vec
from math import tau
import os

def fourier_transform_drawing(input_image, output_animation, frames, coefficients):
    # Convert input_image to an OpenCV image
    input_image = np.array(input_image)
    img = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)

    # processing
    imgray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(imgray, (7, 7), 0)

    (T, 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([len(c) for c in contours])
    verts = [tuple(coord) for coord in contours[largest_contour_idx].squeeze()]

    xs, ys = zip(*verts)
    xs = np.asarray(xs) - np.mean(xs)
    ys = - np.asarray(ys) + np.mean(ys)
    t_list = np.linspace(0, tau, len(xs))

    # Compute the Fourier coefficients
    def f(t, t_list, xs, ys):
        return np.interp(t, t_list, xs + 1j*ys)

    def compute_cn(f, n):
        coef = 1/tau*quad_vec(
            lambda t: f(t, t_list, xs, ys)*np.exp(-n*t*1j), 
            0, 
            tau, 
            limit=100, 
            full_output=False)[0]
        return coef

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

    # animate the drawings
    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(-500, 500)
    ax.set_ylim(-500, 500)
    ax.set_axis_off()
    ax.set_aspect('equal')
    fig.set_size_inches(15, 15)

    draw_x, draw_y = [], []
    
    def animate(i, coefs, time):
        t = time[i]
        coefs = [(c * np.exp(1j*(fr * tau * t)), fr) for c, fr in coefs]
        center = (0, 0)

        for c, _ in coefs:
            r = np.linalg.norm(c)
            theta = np.linspace(0, tau, 80)
            x, y = center[0] + r * np.cos(theta), center[1] + r * np.sin(theta)
            circle_lines[_].set_data([center[0], center[0]+np.real(c)], [center[1], center[1]+np.imag(c)])
            circles[_].set_data(x, y) 
            center = (center[0] + np.real(c), center[1] + np.imag(c))
        
        draw_x.append(center[0])
        draw_y.append(center[1])
        drawing.set_data(draw_x, draw_y)

    drawing_time = 1
    time = np.linspace(0, drawing_time, num=frames)    
    anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, time)) 

    anim.save(output_animation, fps=15)
    plt.close(fig)

    return output_animation

# Gradio interface
interface = gr.Interface(
    fn=fourier_transform_drawing,
    inputs=[
        
        gr.Image(label="Input Image", sources=['upload'], type="pil"),
        gr.Textbox(label="Output Animation Path"),
        gr.Slider(minimum=10, maximum=500, value=300, label="Number of Frames"),
        gr.Slider(minimum=10, maximum=500, value=300, label="Number of Coefficients")
    ],
    outputs="file",
    title="Fourier Transform Drawing",
    description="Upload an image and generate a Fourier Transform drawing animation."
)

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