import argparse
import glob
import os.path

import time
import datetime
from pytz import timezone

import gradio as gr

import pickle
import tqdm
import json

import TMIDIX
from midi_to_colab_audio import midi_to_colab_audio

import copy
from collections import Counter
import random
import statistics

import matplotlib.pyplot as plt

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

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

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

def find_midi(search_string, search_options):
    print('=' * 70)
    print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
    start_time = time.time()
    
    print('=' * 70)
    print('Search string:', search_string)
    print('Search options:', search_options)
    
    print('=' * 70)
    print('Preparing to search...')
    
    #==================================================
    
    random.shuffle(AUX_DATA)
    
    search_data = []
    
    for A in AUX_DATA:
        data = ''
        if 'Titles' in search_options:
            data += A[1] + '\n\n'
        if 'Lyrics' in search_options:
            data += A[2] + '\n\n'
        if 'Summaries' in search_options:
            data += A[3] + '\n\n'

        search_data.append(data)
    
    print('Searching titles...Please wait...')
    
    search_match_data = TMIDIX.ascii_texts_search(search_data, search_string, deterministic_matching = True)
    
    search_match_text = search_match_data[0]
    search_match_ratio = search_match_data[1]
    search_match_index = search_data.index(search_match_text)
    
    print('Done!')
    print('=' * 70)
    print('Search match ratio:', search_match_ratio)
    print('Selected file/title:', AUX_DATA[search_match_index][:2])
    print('=' * 70)

    fn = AUX_DATA[search_match_index][0]
    title = AUX_DATA[search_match_index][1]
    lyric = AUX_DATA[search_match_index][2]
    summary = AUX_DATA[search_match_index][3]
    raw_score = AUX_DATA[search_match_index][4]
    single_track_score_notes = TMIDIX.advanced_score_processor(raw_score, 
                                                               return_score_analysis=False, 
                                                               return_enhanced_score_notes=True)[0]

    print('Sample INTs', raw_score[1][:5])
    print('=' * 70)
 
    x = []
    y = []
    c = []
    
    colors = ['red', 'yellow', 'green', 'cyan',
            'blue', 'pink', 'orange', 'purple',
            'gray', 'white', 'gold', 'silver',
            'lightgreen', 'indigo', 'maroon', 'turquoise']
    
    for s in single_track_score_notes:
        x.append(s[1])
        y.append(s[4])
        c.append(colors[s[3]])

    plt.close()
    plt.figure(figsize=(14,5))
    ax=plt.axes(title=title)
    ax.set_facecolor('black')
    
    plt.scatter(x,y, s=10, c=c)
    plt.xlabel("Time in MIDI ticks")
    plt.ylabel("MIDI Pitch")

    with open(fn+'.mid', 'wb') as f:
        f.write(TMIDIX.score2midi(raw_score))
    
    audio = midi_to_colab_audio(fn+'.mid', 
                        soundfont_path=soundfont_path, 
                        sample_rate=16000, # 44100
                        volume_scale=10,
                        output_for_gradio=True
                        )

    print('Req end time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
    print('-' * 70)
    print('Req execution time:', (time.time() - start_time), 'sec')
    
    yield AUX_DATA[search_match_index][0], str(search_match_ratio), AUX_DATA[search_match_index][1], AUX_DATA[search_match_index][2], AUX_DATA[search_match_index][3],  fn+'.mid', (16000, audio), plt

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

if __name__ == "__main__":

    PDT = timezone('US/Pacific')
    
    print('=' * 70)
    print('App start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
    print('=' * 70)

    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 = "English_Karaoke_Files_Titles_Lyrics_Summaries_Scores_Final.pickle"

    print('Loading meta-data...')
    with open(meta_data_path, 'rb') as f:
        AUX_DATA = pickle.load(f)
    print('Done!')
    
    app = gr.Blocks()
    with app:
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Karaoke MIDI Search</h1>")
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Search and explore 5865 select Karaoke MIDI titles</h1>")
        
        gr.Markdown("![Visitors](https://api.visitorbadge.io/api/visitors?path=asigalov61.Karaoke-MIDI-Search&style=flat)\n\n"
                    "Los Angeles MIDI Dataset Demo\n\n"
                    "Please see [Los Angeles MIDI Dataset](https://github.com/asigalov61/Los-Angeles-MIDI-Dataset) for more information and features\n\n"
                    "[Open In Colab]"
                    "(https://colab.research.google.com/github/asigalov61/Los-Angeles-MIDI-Dataset/blob/main/Los_Angeles_MIDI_Dataset_Search_and_Explore.ipynb)"
                    " for all features\n\n"
                    "Presented lyrics titles, transcriptions, and summaries were auto-generated with [Mistral-7B-Instruct-v0.2](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2) so please keep it in mind"
                    )
        
        search_string = gr.Textbox(label="Enter search prompt here", value="So close, no matter how far\nCouldn't be much more from the heart\nForever trusting who we are\nAnd nothing else matters")
        search_options = gr.CheckboxGroup(["Titles", "Lyrics", "Summaries"], value="Lyrics", label="Search within")
        submit = gr.Button()

        gr.Markdown("# Search results")
        
        output_midi_search_match_ratio = gr.Textbox(label="Output Karaoke MIDI search match ratio")
        output_midi_md5 = gr.Textbox(label="Output Karaoke MIDI md5 hash")
        output_midi_title = gr.Textbox(label="Output Karaoke MIDI title")
        output_midi_summary = gr.Textbox(label="Output Karaoke MIDI summary")
        output_audio = gr.Audio(label="Output Karaoke MIDI audio", format="wav", elem_id="midi_audio")
        output_plot = gr.Plot(label="Output Karaoke MIDI score plot")
        output_midi_lyric = gr.Textbox(label="Output Karaoke MIDI lyric")
        output_midi = gr.File(label="Output Karaoke MIDI original file", file_types=[".mid"])
        
        run_event = submit.click(find_midi, [search_string, search_options],
                                                  [output_midi_md5, output_midi_search_match_ratio, output_midi_title, output_midi_lyric, output_midi_summary, output_midi, output_audio, output_plot])
        
    app.queue(1).launch(server_port=opt.port, share=opt.share, inbrowser=True)