import argparse
import glob
import os.path

import torch
import torch.nn.functional as F

import gradio as gr
import numpy as np
import onnxruntime as rt
import tqdm
import json

from midi_synthesizer import synthesis
import TMIDIX

in_space = os.getenv("SYSTEM") == "spaces"

providers = ['CPUExecutionProvider']

#=================================================================================================

def generate(
            start_tokens,
            seq_len,
            max_seq_len = 2048,
            temperature = 0.9,
            verbose=True,
            return_prime=False,
            ):

    out = torch.LongTensor([start_tokens])

    st = len(start_tokens)

    if verbose:
      print("Generating sequence of max length:", seq_len)

    for s in range(seq_len):
        x = out[:, -max_seq_len:]

        torch_in = x.tolist()[0]

        logits = torch.FloatTensor(session.run(None, {'input': [torch_in]})[0])[:, -1]
        
        filtered_logits = logits

        probs = F.softmax(filtered_logits / temperature, dim=-1)

        sample = torch.multinomial(probs, 1)

        out = torch.cat((out, sample), dim=-1)

        if verbose:
          if s % 32 == 0:
            print(s, '/', seq_len)

    if return_prime:
      return out[:, :]
    
    else:
      return out[:, st:]


#=================================================================================================


def load_javascript(dir="javascript"):
    scripts_list = glob.glob(f"{dir}/*.js")
    javascript = ""
    for path in scripts_list:
        with open(path, "r", encoding="utf8") as jsfile:
            javascript += f"\n<!-- {path} --><script>{jsfile.read()}</script>"
    template_response_ori = gr.routes.templates.TemplateResponse

    def template_response(*args, **kwargs):
        res = template_response_ori(*args, **kwargs)
        res.body = res.body.replace(
            b'</head>', f'{javascript}</head>'.encode("utf8"))
        res.init_headers()
        return res

    gr.routes.templates.TemplateResponse = template_response


class JSMsgReceiver(gr.HTML):

    def __init__(self, **kwargs):
        super().__init__(elem_id="msg_receiver", visible=False, **kwargs)

    def postprocess(self, y):
        if y:
            y = f"<p>{json.dumps(y)}</p>"
        return super().postprocess(y)

    def get_block_name(self) -> str:
        return "html"

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--share", action="store_true", default=False, help="share gradio app")
    parser.add_argument("--port", type=int, default=7860, help="gradio server port")
    parser.add_argument("--max-gen", type=int, default=1024, help="max")
    opt = parser.parse_args()


    providers = ['CPUExecutionProvider']

    session = rt.InferenceSession('Allegro_Music_Transformer_Small_Trained_Model_56000_steps_0.9399_loss_0.7374_acc.onnx', providers=providers)


    app = gr.Blocks()
    with app:
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Midi Composer</h1>")
        gr.Markdown("![Visitors](https://api.visitorbadge.io/api/visitors?path=skytnt.midi-composer&style=flat)\n\n"
                    "Midi event transformer for music generation\n\n"
                    "Demo for [SkyTNT/midi-model](https://github.com/SkyTNT/midi-model)\n\n"
                    "[Open In Colab]"
                    "(https://colab.research.google.com/github/SkyTNT/midi-model/blob/main/demo.ipynb)"
                    " for faster running and longer generation"
                        )
        
        js_msg = JSMsgReceiver()
        tab_select = gr.Variable(value=0)
    
        app.queue(2).launch(server_port=opt.port, share=opt.share, inbrowser=True)