import gradio as gr
import numpy as np
import librosa
import os
import soundfile as sf
import generation_utilities

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# Default song and similarity values
song_default = np.random.choice(["22", "Anti-Hero", "Back-to-december","Blank-Space","Cardigan","Delicate","Lover","Love-Story","Willow","You-Belong-With-Me"])
similarity_default = round(np.random.uniform(0.8, 0.99), 2)

def generate_song(user_id, song_options, similarity):
    # Load songs
    song_list = [librosa.load(os.path.join(os.getcwd(), f"input_songs/{song}.mp3"), sr=22050)[0] for song in song_options]
    
    # Generate spectrogram and song
    spectrogram, generated_song, model_name = generation_utilities.generate_songs(song_list, similarity=similarity, quality=500, merging_quality=100)

    # Save generated song and spectrogram
    sf.write("ui/temp.wav", generated_song, 22050)
    np.save("ui/temp.npy", spectrogram)

    # Return user info, generated song path, and link to rating page
    return {
        "user_id": user_id, 
        "song_list": song_options, 
        "similarity": similarity, 
        "model_name": model_name, 
        "generated_song": "ui/temp.wav", 
        "message": "Song generated! [Click here to go to the rating page](ui/gradio_rating.py)"
    }

# Gradio Interface
with gr.Blocks() as demo:
    user_id = gr.Textbox(label="Enter your user ID")
    song_options = gr.CheckboxGroup(["22", "Anti-Hero", "Back-to-december","Blank-Space","Cardigan","Delicate","Lover","Love-Story","Willow","You-Belong-With-Me"], label="Select songs from library", value=[song_default])
    similarity = gr.Slider(minimum=0.0, maximum=1.0, value=similarity_default, label="Similarity")

    output = gr.JSON(label="Session Info")
    generate_button = gr.Button("Generate Song")
    generate_button.click(fn=generate_song, inputs=[user_id, song_options, similarity], outputs=output)

demo.launch()