import argparse
import glob
import os.path

import gradio as gr

import pickle
import tqdm
import json

import MIDI
from midi_synthesizer import synthesis

import matplotlib.pyplot as plt

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

def run(search_prompt, mid=None, progress=gr.Progress()):
    
    if mid == None:
        
        for m in progress.tqdm(meta_data):
            mid_seq = m[1][17:-1]
            mid_seq_ticks = m[1][16][1]
            mdata = m[1][:16]
            break
       
    elif mid is not None:
        mid_seq = MIDI.midi2score(mid)


    x = []
    y = []
    c = []
    
    colors = ['red', 'yellow', 'green', 'cyan',
            'blue', 'pink', 'orange', 'purple',
            'gray', 'white', 'gold', 'silver',
            'lightgreen', 'indigo', 'maroon', 'turquoise']
    
    for s in [m for m in mid_seq if m[0] == 'note']:
        x.append(s[1] / mid_seq_ticks)
        y.append(s[4])
        c.append(colors[s[3]])

    plt.close()
    plt.figure(figsize=(14,5))
    ax=plt.axes(title='MIDI Search Plot')
    ax.set_facecolor('black')
    
    plt.scatter(x,y, c=c)
    plt.xlabel("Time")
    plt.ylabel("Pitch")
 
    with open(f"output.mid", 'wb') as f:
        f.write(MIDI.score2midi([mid_seq_ticks, mid_seq]))
    audio = synthesis(MIDI.score2opus([mid_seq_ticks, mid_seq]), soundfont_path)
    yield mdata, "output.mid", (44100, audio), plt

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()
    
    soundfont_path = "SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2"
    meta_data_path = "meta-data/LAMD_META_10000.pickle"
    
    models_info = {"generic pretrain model": ["skytnt/midi-model", ""],
                   "j-pop finetune model": ["skytnt/midi-model-ft", "jpop/"],
                   "touhou finetune model": ["skytnt/midi-model-ft", "touhou/"]}


    print('Loading meta-data...')
    with open(meta_data_path, 'rb') as f:
        meta_data = pickle.load(f)
    print('Done!')
    
    app = gr.Blocks()
    with app:
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>MIDI Search</h1>")
        gr.Markdown("![Visitors](https://api.visitorbadge.io/api/visitors?path=asigalov61.MIDI-Search&style=flat)\n\n"
                    "MIDI Search and Explore\n\n"
                    "Demo for [MIDI Search](https://github.com/asigalov61)\n\n"
                    "[Open In Colab]"
                    "(https://colab.research.google.com/github/asigalov61/MIDI-Search/blob/main/demo.ipynb)"
                    " for faster running and longer generation"
                    )
        
        with gr.Tabs():
            with gr.TabItem("instrument prompt") as tab1:
                
                search_prompt = gr.Textbox(label="search prompt")
                
            with gr.TabItem("midi prompt") as tab2:
                input_midi = gr.File(label="input midi", file_types=[".midi", ".mid"], type="binary")

        with gr.Accordion("options", open=False):
 
            input_allow_cc = gr.Checkbox(label="allow midi cc event", value=True)
            
        search_btn = gr.Button("search", variant="primary")
        
        output_audio = gr.Audio(label="output audio", format="mp3", elem_id="midi_audio")
        output_midi = gr.File(label="output midi", file_types=[".mid"])
        output_midi_seq = gr.Textbox(label="output midi metadata")
        output_plot = gr.Plot(label="output midi plot")
        
        run_event = search_btn.click(run, [search_prompt],
                                  [output_midi_seq, output_midi, output_audio, output_plot])
        
    app.queue(1).launch(server_port=opt.port, share=opt.share, inbrowser=True)