import gradio as gr
import torch
import torchaudio
import hydra
from hydra import compose, initialize
import random
import os
from remfx import effects

cfg = None
classifier = None
models = {}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

ALL_EFFECTS = effects.Pedalboard_Effects


def init_hydra():
    global cfg
    initialize(config_path="cfg", job_name="remfx", version_base="2.0")
    cfg = compose(config_name="config", overrides=["+exp=remfx_detect"])


def load_models():
    global classifier
    print("Loading models")
    classifier = hydra.utils.instantiate(cfg.classifier, _convert_="partial")
    ckpt_path = cfg.classifier_ckpt
    state_dict = torch.load(ckpt_path, map_location=device)["state_dict"]
    classifier.load_state_dict(state_dict)
    classifier.to(device)

    for effect in cfg.ckpts:
        model = hydra.utils.instantiate(cfg.ckpts[effect].model, _convert_="partial")
        ckpt_path = cfg.ckpts[effect].ckpt_path
        state_dict = torch.load(ckpt_path, map_location=device)["state_dict"]
        model.load_state_dict(state_dict)
        model.to(device)
        models[effect] = model


def audio_classification(audio_file):
    audio, sr = torchaudio.load(audio_file)
    audio = torchaudio.transforms.Resample(sr, cfg.sample_rate)(audio)
    # Convert to mono
    audio = audio.mean(0, keepdim=True)
    # Add dimension for batch
    audio = audio.unsqueeze(0)
    audio = audio.to(device)

    with torch.no_grad():
        # Classify
        print("Detecting effects")
        labels = torch.tensor(classifier(audio))
        labels_dict = {
            ALL_EFFECTS[i].__name__.replace("RandomPedalboard", ""): labels[i].item()
            for i in range(len(ALL_EFFECTS))
        }
    return labels_dict


def audio_removal(audio_file, labels, threshold):
    audio, sr = torchaudio.load(audio_file)
    audio = torchaudio.transforms.Resample(sr, cfg.sample_rate)(audio)
    # Convert to mono
    audio = audio.mean(0, keepdim=True)
    # Add dimension for batch
    audio = audio.unsqueeze(0)
    audio = audio.to(device)

    label_names = [f"RandomPedalboard{lab['label']}" for lab in labels["confidences"]]
    logits = torch.tensor([lab["confidence"] for lab in labels["confidences"]])
    rem_fx_labels = torch.where(logits > threshold, 1.0, 0.0)
    effects_present = [
        name for name, effect in zip(label_names, rem_fx_labels) if effect == 1.0
    ]
    print("Removing effects:", effects_present)
    # Remove effects
    # Shuffle effects order
    effects_order = cfg.inference_effects_ordering
    random.shuffle(effects_order)
    # Get the correct effect by search for names in effects_order
    effects = [effect for effect in effects_order if effect in effects_present]
    elem = audio
    with torch.no_grad():
        for effect in effects:
            # Sample the model
            elem = models[effect].model.sample(elem)
        output = elem.squeeze(0)
    waveform = gr.make_waveform((cfg.sample_rate, output[0].numpy()))

    return waveform


def ui():
    css = """

    #classifier {
        padding-top: 40px;
    }
    #classifier .output-class {
        display: none;

    }
    """
    with gr.Blocks(css=css) as interface:
        gr.HTML(
            """
                <div style="text-align: center; max-width: 700px; margin: 0 auto;">
                <div
                    style="
                    display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;
                    "
                >
                    <h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
                    RemFx: General Purpose Audio Effect Removal
                    </h1>
                </div> <p style="margin-bottom: 10px; font-size: 94%">
                    <a href="https://arxiv.org/abs/2308.16177">[Paper]</a> <a href="https://csteinmetz1.github.io/RemFX/">[Project
                    page]</a>
                </p>
                </div>
            """
        )
        gr.HTML(
            """
                <div style="text-align: left;"> This is our demo for the paper General Purpose Audio Effect Removal. It uses the RemFX Detect system described in the paper to detect the audio effects that are present and remove them. <br>
                To use the demo, use one of our curated examples or upload your own audio file and click submit. The system will then detect the effects present in the audio remove them if they meet the threshold. </div>
            """
        )
        with gr.Row():
            with gr.Column():
                effected_audio = gr.Audio(
                    source="upload",
                    type="filepath",
                    label="File",
                    interactive=True,
                    elem_id="melody-input",
                )
                submit = gr.Button("Submit")
                threshold = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    step=0.1,
                    value=0.5,
                    label="Detection Threshold",
                )
            with gr.Column():
                classifier = gr.Label(
                    num_top_classes=5, label="Effects Present", elem_id="classifier"
                )
                audio_output = gr.Video(label="Output")
        gr.Examples(
            fn=audio_removal,
            examples=[
                ["./input_examples/guitar.wav"],
                ["./input_examples/vocal.wav"],
                ["./input_examples/bass.wav"],
                ["./input_examples/drums.wav"],
                ["./input_examples/crazy_guitar.wav"],
            ],
            inputs=effected_audio,
        )
        submit.click(
            audio_classification,
            inputs=[effected_audio],
            outputs=[classifier],
            queue=False,
            show_progress=False,
        ).then(
            audio_removal,
            inputs=[effected_audio, classifier, threshold],
            outputs=[audio_output],
        )

        interface.queue().launch()


if __name__ == "__main__":
    init_hydra()
    load_models()
    ui()