import torch
import torchaudio
import gradio as gr
from demucs import pretrained
from demucs.apply import apply_model
from audiotools import AudioSignal
from typing import Dict
from pyharp import *


#DEMUX_MODELS = ["mdx_extra_q", "mdx_extra", "htdemucs", "mdx_q"]

STEM_CHOICES = {
    "Vocals": 3,
    "Drums": 0,
    "Bass": 1,
    "Other": 2,
    "Instrumental (No Vocals)": "instrumental"
}


#models = dict(zip(DEMUX_MODELS, [pretrained.get_model(m) for m in DEMUX_MODELS]))

#for model in models.values():
    #model.eval()

model = pretrained.get_model('mdx_extra_q')


def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> AudioSignal:
    waveform, sr = torchaudio.load(audio_file_path)
    is_mono = waveform.shape[0] == 1
    if is_mono:
        waveform = waveform.repeat(2, 1)

    with torch.no_grad():
        stems_batch = apply_model(
            model,
            waveform.unsqueeze(0),
            overlap=0.2,
            shifts=1,
            split=True,
            progress=True,
            num_workers=4
        )

    stems = stems_batch[0]

    if stem_choice == "Instrumental (No Vocals)":
        stem = stems[0] + stems[1] + stems[2]
    else:
        stem_index = STEM_CHOICES[stem_choice]
        stem = stems[stem_index]

    if is_mono:
        stem = stem.mean(dim=0, keepdim=True)

    return AudioSignal(stem.cpu().numpy().astype('float32'), sample_rate=sr)

# Gradio Callback Function

def process_fn_stem(audio_file_path: str, stem_choice: str):
    """
    PyHARP process function:
      - Separates the chosen stem using Demucs.
      - Saves the stem as a .wav file.
    """
    stem_signal = separate_stem(audio_file_path, model_name='', stem_choice=stem_choice)
    stem_path = save_audio(stem_signal, f"{stem_choice.lower().replace(' ', '_')}.wav")
    return stem_path, LabelList(labels=[])


# Model Card 
model_card = ModelCard(
    name="Demucs Stem Separator",
    description="Uses Demucs to separate a music track into a selected stem.",
    author="Alexandre Défossez, Nicolas Usunier, Léon Bottou, Francis Bach",
    tags=["demucs", "source-separation", "pyharp", "stems"]
)

# Gradio UI
with gr.Blocks() as demo:

    #dropdown_model = gr.Dropdown(
    #    label="Demucs Model",
    #    choices=DEMUX_MODELS,
    #    value="mdx_extra_q"
    #)

    dropdown_stem = gr.Dropdown(
        label="Stem to Separate",
        choices=list(STEM_CHOICES.keys()),
        value="Vocals"
    )

    app = build_endpoint(
        model_card=model_card,
        components=[dropdown_stem],
        process_fn=process_fn_stem
    )

demo.queue()
demo.launch(show_error=True)