import streamlit as st
import librosa
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import tempfile
import os

# Import custom visualizers
from felguk import felguk_viz
from result_video_audio_vizualizer import result_video_viz
from piepy import piepy_viz
from pyz import pyz_viz

# List of visualizers
VISUALIZERS = {
    "Waveform": lambda y, sr: librosa.display.waveshow(y, sr=sr),
    "Spectrogram": lambda y, sr: librosa.display.specshow(librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max), sr=sr, x_axis='time', y_axis='log'),
    "Mel Spectrogram": lambda y, sr: librosa.display.specshow(librosa.feature.melspectrogram(y=y, sr=sr), sr=sr, x_axis='time', y_axis='mel'),
    "Chromagram": lambda y, sr: librosa.display.specshow(librosa.feature.chroma_stft(y=y, sr=sr), sr=sr, x_axis='time', y_axis='chroma'),
    "Felguk Visualizer": felguk_viz,
    "Result Video Visualizer": result_video_viz,
    "Piepy Visualizer": piepy_viz,
    "Pyz Visualizer": pyz_viz,
    "Tempogram": lambda y, sr: librosa.display.specshow(librosa.feature.tempogram(y=y, sr=sr)),
    "Spectral Contrast": lambda y, sr: librosa.display.specshow(librosa.feature.spectral_contrast(y=y, sr=sr), sr=sr, x_axis='time'),
}

# Function to generate a video
def generate_video(y, sr, visualizer):
    fig, ax = plt.subplots(figsize=(10, 4))
    def animate(i):
        ax.clear()
        visualizer(y[:i * 1000], sr)  # Update visualization for each frame
        ax.set_title(f"Frame {i}")
    anim = FuncAnimation(fig, animate, frames=len(y) // 1000, interval=50)
    
    # Save the animation as a video file
    video_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    anim.save(video_file.name, writer="ffmpeg", fps=30)
    plt.close()
    return video_file.name

# Streamlit app
st.title("Audio Visualizer")
st.write("Upload an audio file to visualize it in different ways.")

# Upload audio file
audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "ogg"])

if audio_file is not None:
    st.audio(audio_file, format='audio/wav')
    y, sr = librosa.load(audio_file, sr=None)

    # Select visualizer
    selected_viz = st.selectbox("Select a visualizer", list(VISUALIZERS.keys()))

    # Plot the selected visualizer
    st.write(f"### {selected_viz}")
    plt.figure(figsize=(10, 4))
    VISUALIZERS[selected_viz](y, sr)
    plt.colorbar(format="%+2.0f dB")
    st.pyplot(plt)

    # Generate and return video
    if st.button("Generate Video"):
        video_path = generate_video(y, sr, VISUALIZERS[selected_viz])
        st.video(video_path)
        with open(video_path, "rb") as file:
            st.download_button(
                label="Download Video",
                data=file,
                file_name="audio_visualization.mp4",
                mime="video/mp4"
            )
        os.unlink(video_path)  # Clean up the temporary file