import streamlit as st
import os
import numpy as np
from inference import generate_drum_kit
from audio_utils import play_audio
from fx import get_fx

# Streamlit UI
st.title("semantic spaces: kit generator")
st.subheader("generate drum kits and audio effects with text prompts")
st.write("uses publicly available samples from [freesound](https://zenodo.org/records/4687854) and [CLAP embeddings](https://github.com/LAION-AI/CLAP) for text-based querying")
st.write("hint: turn audio effects on! try weird prompts!")

with st.container(border=True):
    # User Inputs
    prompt = st.text_input("Describe your drum kit:", "warm vintage organic percussion")
    kit_size = st.slider("Number of sounds per instrument:", 1, 10, 4)
    use_fx = st.toggle("Apply audio effects?", value=True)
    if use_fx:
        if st.toggle("Use a different prompt for audio effects?", value=True):
            fx_prompt = st.text_input("Describe your desired FX tone:", "soft and ethereal space")
        else:
            fx_prompt = prompt

    # Run the inference
    if st.button("Generate Drum Kit"):
        drum_kit = generate_drum_kit(prompt, kit_size)
        st.session_state["dry_kit"] = drum_kit
        if use_fx:
            drum_kit, fx_params = get_fx(drum_kit, fx_prompt)
            st.session_state["fx_params"] = fx_params
        st.session_state["drum_kit"] = drum_kit  # Store results

# Display results
if "drum_kit" in st.session_state:
    drum_kit = st.session_state["drum_kit"]
    st.subheader("Generated Drum Kit")

    for instrument, sounds in drum_kit.items():
        st.write(f"**{instrument}**")
        cols = st.columns(len(sounds))

        for i, sound_file in enumerate(sounds):
            with cols[i]:
                if st.button(f"▶️ {os.path.basename(sound_file)}", key=sound_file):
                    play_audio(sound_file)
    
    if st.toggle("Show parameters?"):
        if "fx_params" in st.session_state:
            st.subheader("FX Parameters")
            st.write(st.session_state["fx_params"])
        if "dry_kit" in st.session_state:
            st.write(st.session_state["dry_kit"])