import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.colors import to_hex
import librosa
import tempfile

def check_rgba(string:str):
    if '#' in string:
        return string
    else:
        clean_str_list = string.split("(")[1].split(")")[0].split(",")
        clean_list = [float(value.strip()) / 255 if idx != 3 else float(value.strip()) / 1 for idx, value in enumerate(clean_str_list)]
        hex = to_hex(clean_list, keep_alpha=True)
        return hex

def extract_waveform_animation(audio_file, window_seconds, waveform_color, background_color):
    y, sr = librosa.load(audio_file, sr=None)
    duration = librosa.get_duration(y=y, sr=sr)
    FPS = 1
    fig, ax = plt.subplots()
    line, = ax.plot([], [], lw=2, color=check_rgba(waveform_color))
    window_length = int(window_seconds * sr)
    x_vals = np.linspace(0, duration, num=len(y))
    ax.set_axis_off()
    bg_color = check_rgba(background_color)
    fig.set_facecolor(bg_color)
    
    def init():
        ax.set_xlim(0, window_seconds)
        ax.set_ylim(np.min(y), np.max(y))  # Reduced max for visibility
        return line,

    def update(frame):
        # Get current window
        start = frame * sr
        end = start + window_length
        window = y[start:end]
        
        # Update x and y limits
        ax.set_xlim(frame, frame + window_seconds)
        
        # Update line data
        line.set_data(x_vals[start:end], window)
        return line,

    total_frames = int(duration) * FPS
    ani = FuncAnimation(fig, update, frames=range(total_frames), 
                       init_func=init, interval=window_seconds, blit=False)
    
    with tempfile.NamedTemporaryFile(delete=False, suffix='.gif') as tmpfile:
        ani.save(tmpfile.name, writer='ffmpeg', fps=FPS)
        video_path = tmpfile.name
    
    return video_path

# Modified interface with window controls
iface = gr.Interface(
    fn=extract_waveform_animation,
    inputs=[
        gr.Audio(type="filepath"),
        gr.Slider(1, 10, value=5, step=1, label="Window Size (seconds)"),
        gr.ColorPicker(label="Waveform color",  value="#0021e4"),
        gr.ColorPicker(label="Background color",  value="#00FFFF00")
    ],
    outputs=gr.Image(),
    description="Scroll through audio waveform with a moving window."
)

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