import gradio as gr
import os
from transformer_wrapper import TransformerWrapper
from omegaconf import OmegaConf


def get_file_content_as_string(path):
    return open(path, "r", encoding="utf-8").read()


def model_load():
    config = OmegaConf.load("config.yaml")
    wrapper = TransformerWrapper(config)
    wrapper = wrapper.load_from_checkpoint(
        "https://huggingface.co/sweetcocoa/pop2piano/resolve/main/model-1999-val_0.67311615.ckpt",
        config=config,
        map_location="cpu",
    )
    model_id = "dpipqxiy"
    wrapper.eval()
    return wrapper, model_id, config


wrapper, model_id, config = model_load()
composers = list(config.composer_to_feature_token.keys())
dest_dir = "ytsamples"
os.makedirs(dest_dir, exist_ok=True)


def inference(file_up, composer):

    midi, arranger, mix_path, midi_path = wrapper.generate(
        audio_path=file_up,
        composer=composer,
        model=model_id,
        ignore_duplicate=True,
        show_plot=False,
        save_midi=True,
        save_mix=True,
    )

    return mix_path


block = gr.Blocks()


with block:
    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;">
                  Pop2piano
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%">
                A demo for Pop2Piano:Pop Audio-based Piano Cover Generation. Please select the composer and upload the pop audio to submit.
              </p>
            </div>
        """
    )
    with gr.Group():
        with gr.Box():
            with gr.Row().style(mobile_collapse=False, equal_height=True):
                file_up = gr.Audio(label="Upload an audio", type="filepath")
                composer = gr.Dropdown(label="Arranger", choices=composers, value="composer1")
                btn = gr.Button("Convert")
        out = gr.Audio(label="Output")

        btn.click(inference, inputs=[file_up, composer], outputs=out)
        gr.HTML(
            """
        <div class="footer">
                    <p><a href="http://sweetcocoa.github.io/pop2piano_samples" style="text-decoration: underline;" target="_blank">Project Page</a>
                    </p>
        </div>
        """
        )

block.launch(debug=True)