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

def extract_waveform_animation(audio_file, window_seconds=5):
    y, sr = librosa.load(audio_file, sr=None)
    duration = librosa.get_duration(y=y, sr=sr)
    FPS = 30
    fig, ax = plt.subplots()
    line, = ax.plot([], [], lw=2)
    window_length = int(window_seconds * sr / FPS)
    
    # Initialize with first window
    first_window = y[:window_length]
    x_vals = np.linspace(0, duration, num=len(y))
    ax.set_axis_off()
    fig.set_facecolor("black")
    
    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='.mp4') 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)")
    ],
    outputs=gr.Video(),
    description="Scroll through audio waveform with a moving window."
)

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