import re
import os
import time
import torch
import shutil
import argparse
import gradio as gr
from utils import *
from config import *
from convert import *
from transformers import GPT2Config
import warnings

warnings.filterwarnings("ignore")


def get_args(parser):
    parser.add_argument(
        "-num_tunes",
        type=int,
        default=1,
        help="the number of independently computed returned tunes",
    )
    parser.add_argument(
        "-max_patch",
        type=int,
        default=128,
        help="integer to define the maximum length in tokens of each tune",
    )
    parser.add_argument(
        "-top_p",
        type=float,
        default=0.8,
        help="float to define the tokens that are within the sample operation of text generation",
    )
    parser.add_argument(
        "-top_k",
        type=int,
        default=8,
        help="integer to define the tokens that are within the sample operation of text generation",
    )
    parser.add_argument(
        "-temperature",
        type=float,
        default=1.2,
        help="the temperature of the sampling operation",
    )
    parser.add_argument("-seed", type=int, default=None, help="seed for randomstate")
    parser.add_argument(
        "-show_control_code",
        type=bool,
        default=True,
        help="whether to show control code",
    )
    args = parser.parse_args()

    return args


def generate_abc(args, region):
    patchilizer = Patchilizer()

    patch_config = GPT2Config(
        num_hidden_layers=PATCH_NUM_LAYERS,
        max_length=PATCH_LENGTH,
        max_position_embeddings=PATCH_LENGTH,
        vocab_size=1,
    )

    char_config = GPT2Config(
        num_hidden_layers=CHAR_NUM_LAYERS,
        max_length=PATCH_SIZE,
        max_position_embeddings=PATCH_SIZE,
        vocab_size=128,
    )

    model = TunesFormer(patch_config, char_config, share_weights=SHARE_WEIGHTS)

    filename = WEIGHT_PATH

    if os.path.exists(filename):
        print(f"Weights already exist at '{filename}'. Loading...")

    else:
        download()

    checkpoint = torch.load(filename, map_location=torch.device("cpu"))
    model.load_state_dict(checkpoint["model"])
    model = model.to(device)
    model.eval()

    prompt = template(region)

    tunes = ""
    num_tunes = args.num_tunes
    max_patch = args.max_patch
    top_p = args.top_p
    top_k = args.top_k
    temperature = args.temperature
    seed = args.seed
    show_control_code = args.show_control_code

    print(" HYPERPARAMETERS ".center(60, "#"), "\n")
    args = vars(args)

    for key in args.keys():
        print(f"{key}: {str(args[key])}")

    print("\n", " OUTPUT TUNES ".center(60, "#"))

    start_time = time.time()

    for i in range(num_tunes):
        title_artist = f"T:{region} Fragment\nC:Generated by AI\n"
        tune = f"X:{str(i + 1)}\n{title_artist + prompt}"
        lines = re.split(r"(\n)", tune)
        tune = ""
        skip = False
        for line in lines:
            if show_control_code or line[:2] not in ["S:", "B:", "E:"]:
                if not skip:
                    print(line, end="")
                    tune += line

                skip = False

            else:
                skip = True

        input_patches = torch.tensor(
            [patchilizer.encode(prompt, add_special_patches=True)[:-1]], device=device
        )

        if tune == "":
            tokens = None

        else:
            prefix = patchilizer.decode(input_patches[0])
            remaining_tokens = prompt[len(prefix) :]
            tokens = torch.tensor(
                [patchilizer.bos_token_id] + [ord(c) for c in remaining_tokens],
                device=device,
            )

        while input_patches.shape[1] < max_patch:
            predicted_patch, seed = model.generate(
                input_patches,
                tokens,
                top_p=top_p,
                top_k=top_k,
                temperature=temperature,
                seed=seed,
            )
            tokens = None

            if predicted_patch[0] != patchilizer.eos_token_id:
                next_bar = patchilizer.decode([predicted_patch])

                if show_control_code or next_bar[:2] not in ["S:", "B:", "E:"]:
                    print(next_bar, end="")
                    tune += next_bar

                if next_bar == "":
                    break

                next_bar = remaining_tokens + next_bar
                remaining_tokens = ""

                predicted_patch = torch.tensor(
                    patchilizer.bar2patch(next_bar), device=device
                ).unsqueeze(0)

                input_patches = torch.cat(
                    [input_patches, predicted_patch.unsqueeze(0)], dim=1
                )

            else:
                break

        tunes += f"{tune}\n\n"
        print("\n")

    print("Generation time: {:.2f} seconds".format(time.time() - start_time))
    os.makedirs("./tmp", exist_ok=True)
    timestamp = time.strftime("%a_%d_%b_%Y_%H_%M_%S", time.localtime())
    out_midi = abc_to_midi(tunes, f"./tmp/[{region}]{timestamp}.mid")
    out_xml = abc_to_musicxml(tunes, f"./tmp/[{region}]{timestamp}.musicxml")
    out_mxl = musicxml_to_mxl(f"./tmp/[{region}]{timestamp}.musicxml")
    pdf_file, jpg_file = mxl2jpg(out_mxl)
    wav_file = midi2wav(out_midi)

    return tunes, out_midi, pdf_file, out_xml, out_mxl, jpg_file, wav_file


def inference(region):
    if os.path.exists("./tmp"):
        shutil.rmtree("./tmp")

    parser = argparse.ArgumentParser()
    args = get_args(parser)
    return generate_abc(args, region)


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            region_opt = gr.Dropdown(
                choices=["Mondstadt", "Liyue", "Inazuma", "Sumeru", "Fontaine"],
                value="Mondstadt",
                label="Region genre",
            )
            gen_btn = gr.Button("Generate")

        with gr.Column():
            wav_output = gr.Audio(label="Audio", type="filepath")
            dld_midi = gr.components.File(label="Download MIDI")
            pdf_score = gr.components.File(label="Download PDF score")
            dld_xml = gr.components.File(label="Download MusicXML")
            dld_mxl = gr.components.File(label="Download MXL")
            abc_output = gr.Textbox(label="abc score", show_copy_button=True)
            img_score = gr.Image(label="Staff", type="filepath")

    gen_btn.click(
        inference,
        inputs=region_opt,
        outputs=[
            abc_output,
            dld_midi,
            pdf_score,
            dld_xml,
            dld_mxl,
            img_score,
            wav_output,
        ],
    )

demo.launch(share=True)