import streamlit as st
import torch
import os
import librosa
import librosa.display
import matplotlib.pyplot as plt
from audiosr import build_model, super_resolution, save_wave
import tempfile
import numpy as np

# Set MPS device if available (for Mac M-Series GPUs)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Title and Description
st.title("AudioSR: Versatile Audio Super-Resolution")
st.write("""
Upload your low-resolution audio files, and AudioSR will enhance them to high fidelity!
Supports all types of audio (music, speech, sound effects, etc.) with arbitrary sampling rates.
""")

# Upload audio file
uploaded_file = st.file_uploader("Upload an audio file (WAV format)", type=["wav"])

# Model Parameters
st.sidebar.title("Model Parameters")
model_name = st.sidebar.selectbox("Select Model", ["basic", "speech"], index=0)
ddim_steps = st.sidebar.slider("DDIM Steps", min_value=10, max_value=100, value=50)
guidance_scale = st.sidebar.slider("Guidance Scale", min_value=1.0, max_value=10.0, value=3.5)
random_seed = st.sidebar.number_input("Random Seed", min_value=0, value=42, step=1)
latent_t_per_second = 12.8

# Helper function to plot spectrogram
def plot_spectrogram(audio_path, title):
    y, sr = librosa.load(audio_path, sr=None)
    S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128, fmax=sr // 2)
    S_dB = librosa.power_to_db(S, ref=np.max)
    
    plt.figure(figsize=(10, 4))
    librosa.display.specshow(S_dB, sr=sr, x_axis='time', y_axis='mel', fmax=sr // 2, cmap='viridis')
    plt.colorbar(format='%+2.0f dB')
    plt.title(title)
    plt.tight_layout()
    return plt

# Process Button
if uploaded_file and st.button("Enhance Audio"):
    st.write("Processing audio...")

    # Create temp directory for saving files
    with tempfile.TemporaryDirectory() as tmp_dir:
        input_path = os.path.join(tmp_dir, "input.wav")
        output_path = os.path.join(tmp_dir, "output.wav")

        # Save uploaded file locally
        with open(input_path, "wb") as f:
            f.write(uploaded_file.read())

        # Plot input spectrogram
        st.write("Input Audio Spectrogram:")
        input_spectrogram = plot_spectrogram(input_path, title="Input Audio Spectrogram")
        st.pyplot(input_spectrogram)

        # Build and load the model
        audiosr = build_model(model_name=model_name, device=device)

        # Perform super-resolution
        waveform = super_resolution(
            audiosr,
            input_path,
            seed=random_seed,
            guidance_scale=guidance_scale,
            ddim_steps=ddim_steps,
            latent_t_per_second=latent_t_per_second,
        )

        # Save enhanced audio
        save_wave(waveform, inputpath=input_path, savepath=tmp_dir, name="output", samplerate=48000)

        # Plot output spectrogram
        st.write("Enhanced Audio Spectrogram:")
        output_spectrogram = plot_spectrogram(output_path, title="Enhanced Audio Spectrogram")
        st.pyplot(output_spectrogram)

        # Display audio players and download link
        st.audio(input_path, format="audio/wav")
        st.write("Original Audio:")
        
        st.audio(output_path, format="audio/wav")
        st.write("Enhanced Audio:")
        st.download_button("Download Enhanced Audio", data=open(output_path, "rb").read(), file_name="enhanced_audio.wav")

# Footer
st.write("Built with [Streamlit](https://streamlit.io) and [AudioSR](https://audioldm.github.io/audiosr)")