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


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

    # processing
    # resize the image to a smaller size for faster processing
    dim = (img_size, img_size)
    img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
    
    imgray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(imgray, (5, 5), 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]

    verts = [tuple(coord) for coord in contours[largest_contour_idx].squeeze()]

    xs, ys = zip(*verts)
    xs, ys = np.asarray(xs), np.asarray(ys)

    # calculate the range of xs and ys
    x_range = np.max(xs) - np.min(xs)
    y_range = np.max(ys) - np.min(ys)
    
    # determine the scale factors
    desired_range = 400
    scale_x = desired_range / x_range
    scale_y = desired_range / y_range
    
    # apply scaling
    # ys needs to be flipped vertically
    xs = (xs - np.mean(xs)) * scale_x
    ys = (-ys + np.mean(ys)) * scale_y

    # compute the Fourier coefficients
    num_points = 1000  # how many points to use for numerical integration
    t_values = np.linspace(0, tau, num_points)
    t_list = np.linspace(0, tau, len(xs))

    def compute_cn(n, t_list, xs, ys):
        """
        Integrate the contour along axis (-1) using the composite trapezoidal rule.
        https://numpy.org/doc/stable/reference/generated/numpy.trapz.html#r7aa6c77779c0-2
        """
        f_exp = np.interp(t, t_list, xs + 1j * ys) * np.exp(-n * t_values * 1j)
        coef = np.trapz(f_exp, t_values) / tau
        return coef

    N = coefficients
    coefs = [(compute_cn(0, t_list, xs, ys), 0)] + [(compute_cn(j, t_list, xs, ys), 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(-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 = [], []
    
    def animate(i, coefs, time):
        t = time[i]
        center = (0, 0)
        theta = np.linspace(0, tau, 80)
        for c, fr in coefs:
            c = c * np.exp(1j*(fr * tau * t)
            r = np.linalg.norm(c)
            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[:i+1], draw_y[:i+1])

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

    # save the animation as an MP4 file
    output_animation = "output.mp4"
    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.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", precision=0)
    ],
    outputs=gr.Video(),
    title="Fourier Transform Drawing",
    description="Upload an image and generate a Fourier Transform drawing animation. You can find out more about the project here : https://github.com/staghado/fourier-draw",
    examples=[["Fourier2.jpg", 100, 100, 224], ["Luffy.png", 150, 200, 224]]
)

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