import gradio as gr
from pathlib import Path

# Sample data for the demo
samples = {"drums": [], "bass": []}

audio_dir: Path = Path("./audio")
bass_dir: Path = audio_dir / "bass"
drums_dir: Path = audio_dir / "drums"

for instrument in samples.keys():
    for subdir in (audio_dir / instrument.lower()).iterdir():
        context = str(sorted(list(subdir.glob("*context.wav")))[0])
        gen = str(sorted(list(subdir.glob("*gen.wav")))[0])
        mix = str(sorted(list(subdir.glob("*mix.wav")))[0])
        samples[instrument].append({
            "context": context,
            "generated": gen,
            "mix": mix
        })

with gr.Blocks(theme=gr.themes.Soft(), css_paths=["styles.css"]) as demo:
    # Header and description
    gr.HTML("<h1>STAGE: demo samples</h1>")
    gr.HTML('''
      <div class="centered-text">
        Listen to audio samples from STAGE, our <i>Single-stem Accompaniment Generation Model</i>.<br>
        All generated samples are straight from the model, with no post-processing.<br>
        All contexts are taken from an isolated test dataset.
        <br>
        <a href="https://github.com/yourname/yourproject" target="_blank">GitHub</a> |
        <a href="https://arxiv.org/abs/yourpaperid" target="_blank">arXiv</a>
      </div>
    ''')

    # Tabs for Bass and Drums samples
    with gr.Tabs(elem_classes="tabs-container"):

        for instrument in samples.keys():

            with gr.Tab(instrument.capitalize(), elem_classes="tab-content"):
                for i, s in enumerate(samples[instrument], 1):
                    with gr.Group(elem_classes="audio-section"):
                        gr.HTML(f"Sample {i}", elem_classes="sample-title")
                        with gr.Row(elem_classes="audio-row"):
                            for name, audiopath in s.items():
                                gr.Audio(audiopath,
                                         label=name.capitalize(),
                                         type="filepath",
                                         elem_classes="audio-player")
demo.launch()