import os
import re
import random
from scipy.io.wavfile import write
import gradio as gr
from Applio import *

demucs_models = [
    'htdemucs_ft.yaml', 
    'htdemucs.yaml',
    'hdemucs_mmi.yaml',
]

output_format = [
    'wav',
    'flac',
    'mp3',
]


demucs_overlap_values = [
    '0.25',
    '0.50',
    '0.75',
    '0.99',
]



def demucs_def(demucs_audio, demucs_model, demucs_output_format, demucs_shifts, demucs_overlap):
  files_list = []
  files_list.clear()
  directory = "./outputs"
  random_id = str(random.randint(10000, 99999))
  pattern = f"{random_id}"
  os.makedirs("outputs", exist_ok=True)
  write(f'{random_id}.wav', demucs_audio[0], demucs_audio[1])
  prompt = f"audio-separator {random_id}.wav --model_filename {demucs_model} --output_dir=./outputs --output_format={demucs_output_format} --normalization=0.9 --demucs_shifts={demucs_shifts} --demucs_overlap={demucs_overlap}"

  os.system(prompt)

  for file in os.listdir(directory):
    if re.search(pattern, file):
      files_list.append(os.path.join(directory, file))

  stem1_file = files_list[0]
  stem2_file = files_list[1]
  stem3_file = files_list[2]
  stem4_file = files_list[3]


with gr.Blocks(title="Demucs Esparation", theme=applio) as demo:
    gr.Markdown("# DEMUCS ESPRATION UVR")
    with gr.TabItem("main settings"):
        demucs_model = gr.Dropdown(
            label = "Select the Model",
            choices = demucs_models,
            interactive = True,
        )
        emucs_output_format = gr.Dropdown(
            label = "Select the Output Format",
            choices = output_format,
            interactive = True,
        )
        demucs_shifts = gr.Slider(
            minimum = 1,
            maximum = 20,
            step = 1,
            label = "Shifts",
            info = "Number of predictions with random shifts, higher = slower but better quality.",
            value = 2,
            interactive = True
        )
        demucs_overlap = gr.Dropdown(
            label = "Overlap",
            choices = demucs_overlap_values,
            value = demucs_overlap_values[0],
            interactive = True
        )
        demucs_audio = gr.Audio(
            label = "Input Audio",
            type = "numpy",
            interactive = True
        ),
        demucs_button = gr.Button("Separate!", variant = "primary")
        demucs_stem1 = gr.Audio(
            show_download_button = True,
            interactive = False,
            type = "filepath",
            label = "Stem 1"
        )
        demucs_stem2 = gr.Audio(
            show_download_button = True,
            interactive = False,
            type = "filepath",
            label = "Stem 2"
        )
        demucs_stem3 = gr.Audio(
            show_download_button = True,
            interactive = False,
            type = "filepath",
            label = "Stem 3"
        )
        demucs_stem4 = gr.Audio(
            show_download_button = True,
            interactive = False,
            type = "filepath",
            label = "Stem 4"
        )

demo.launch()