# =================================================================================================
# https://huggingface.co/spaces/asigalov61/Monophonic-MIDI-Melody-Harmonizer
# =================================================================================================

import os
import time as reqtime
import datetime
from pytz import timezone

import gradio as gr

import os
import random
from tqdm import tqdm

import TMIDIX
import HaystackSearch

from midi_to_colab_audio import midi_to_colab_audio

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

def Generate_Chords_Progression(minimum_song_length_in_chords_chunks,
                                chords_chunks_memory_ratio,
                                chord_time_step,
                                merge_chords_notes,
                                melody_MIDI_patch_number,
                                chords_progression_MIDI_patch_number,
                                base_MIDI_patch_number
                               ):

    print('=' * 70)
    print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
    start_time = reqtime.time()

    print('=' * 70)
    print('Requested settings:')
    print('Minimum song length in chords chunks:', minimum_song_length_in_chords_chunks)
    print('Chords chunks memory ratio:', chords_chunks_memory_ratio)
    print('Chord time step:', chord_time_step)
    print('Merge chords notes max time:', merge_chords_notes)
    print('Melody MIDI patch number:', melody_MIDI_patch_number)
    print('Chords progression MIDI patch number:', chords_progression_MIDI_patch_number)
    print('Base MIDI patch number:', base_MIDI_patch_number)
    print('=' * 70)
    
    #==================================================================

    print('=' * 70)
    print('Pitches Chords Progressions Generator')
    print('=' * 70)

    print('=' * 70)
    print('Chunk-by-chunk generation')
    print('=' * 70)
    print('Generating...')
    print('=' * 70)
    
    matching_long_chords_chunks = []
    
    ridx = random.randint(0, len(all_long_chords_tokens_chunks)-1)
    
    matching_long_chords_chunks.append(ridx)
    
    max_song_len = 0
    
    tries = 0
    
    while len(matching_long_chords_chunks) < minimum_song_length_in_chords_chunks:

        matching_long_chords_chunks = []
    
        ridx = random.randint(0, len(all_long_chords_tokens_chunks)-1)
    
        matching_long_chords_chunks.append(ridx)
        seen = [ridx]
        gseen = [ridx]
    
        for a in range(minimum_song_length_in_chords_chunks * 10):
    
          if not matching_long_chords_chunks:
            break
    
          if len(matching_long_chords_chunks) > minimum_song_length_in_chords_chunks:
            break
    
          schunk = all_long_chords_tokens_chunks[matching_long_chords_chunks[-1]]
          trg_long_chunk = np.array(schunk[-chunk_size:])
          idxs = np.where((src_long_chunks == trg_long_chunk).all(axis=1))[0].tolist()
    
          if len(idxs) > 1:
    
            random.shuffle(idxs)
    
            eidxs = [i for i in idxs if i not in seen]
    
            if eidxs:
              eidx = eidxs[0]
              matching_long_chords_chunks.append(eidx)
              seen.append(eidx)
              gseen.append(eidx)
    
              if 0 < chords_chunks_memory_ratio < 1:
                seen = random.choices(gseen, k=math.ceil(len(gseen) * chords_chunks_memory_ratio))
              elif chords_chunks_memory_ratio == 0:
                seen = []
    
            else:
                gseen.pop()
                matching_long_chords_chunks.pop()
    
          else:
            gseen.pop()
            matching_long_chords_chunks.pop()
    
    
        if len(matching_long_chords_chunks) > max_song_len:
          print('Current song length:', len(matching_long_chords_chunks), 'chords chunks')
          print('=' * 70)
          final_song = matching_long_chords_chunks
    
        max_song_len = max(max_song_len, len(matching_long_chords_chunks))
    
        tries += 1
    
        if tries % 500 == 0:
          print('Number of passed tries:', tries)
          print('=' * 70)

    if len(matching_long_chords_chunks) > max_song_len:
      print('Current song length:', len(matching_long_chords_chunks), 'chords chunks')
      print('=' * 70)
      final_song = matching_long_chords_chunks
    
    f_song = []
    
    for mat in final_song:
      f_song.extend(all_long_good_chords_chunks[mat][:-chunk_size])
    f_song.extend(all_long_good_chords_chunks[mat][-chunk_size:])
    
    print('Generated final song after', tries, 'tries with', len(final_song), 'chords chunks and', len(f_song), 'chords')
    print('=' * 70)

    print('Done!')
    print('=' * 70)
    
    #===============================================================================
    
    print('Rendering results...')
    print('=' * 70)

    output_score = []
    
    time = 0
    
    patches = [0] * 16
    patches[0] = chords_progression_MIDI_patch_number
    
    if base_MIDI_patch_number > -1:
      patches[2] = base_MIDI_patch_number
    
    if melody_MIDI_patch_number > -1:
      patches[3] = melody_MIDI_patch_number
    
    chords_labels = []
    
    for i, s in enumerate(f_song):
    
      time += chord_time_step
    
      dur = chord_time_step
    
      chord_str = str(i+1)
    
      for t in sorted(set([t % 12 for t in s])):
        chord_str += '-' + str(t)
    
      chords_labels.append(['text_event', time, chord_str])
    
      for p in s:
        output_score.append(['note', time, dur, 0, p, max(40, p), chords_progression_MIDI_patch_number])
                    
      if base_MIDI_patch_number > -1:
        output_score.append(['note', time, dur, 2, (s[-1] %  12)+24, 120-(s[-1] %  12), base_MIDI_patch_number])
    
    if melody_MIDI_patch_number > -1:
      output_score = TMIDIX.add_melody_to_enhanced_score_notes(output_score, 
                                                               melody_patch=melody_MIDI_patch_number,
                                                               melody_notes_max_duration=max(merge_chords_notes, chord_time_step)
                                                              )

    if merge_chords_notes > 0:
        escore_matrix = TMIDIX.escore_notes_to_escore_matrix(output_score)
        output_score = TMIDIX.escore_matrix_to_merged_escore_notes(escore_matrix, max_note_duration=merge_chords_notes)
        
    midi_score = sorted(chords_labels + output_score, key=lambda x: x[1])

    fn1 = "Pitches-Chords-Progression-Composition"
    
    detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(midi_score,
                                                              output_signature = 'Pitches Chords Progression',
                                                              output_file_name = fn1,
                                                              track_name='Project Los Angeles',
                                                              list_of_MIDI_patches=patches
                                                              )
    
    new_fn = fn1+'.mid'
            
    
    audio = midi_to_colab_audio(new_fn, 
                        soundfont_path=soundfont,
                        sample_rate=16000,
                        volume_scale=10,
                        output_for_gradio=True
                        )
    
    #========================================================

    output_midi_title = str(fn1)
    output_midi = str(new_fn)
    output_audio = (16000, audio)
    
    output_plot = TMIDIX.plot_ms_SONG(output_score, plot_title=output_midi, return_plt=True)
    
    print('Done!')
    print('=' * 70)
    
    #========================================================

    print('Generated chords progression info and stats:')
    print('=' * 70)

    chords_progression_summary_string = '=' * 70
    chords_progression_summary_string += '\n'

    all_song_chords = []
    
    for pc in f_song:
      tones_chord = tuple(sorted(set([p % 12 for p in pc])))
      all_song_chords.append([pc, tones_chord])
    
    print('=' * 70)
    print('Total number of chords:', len(all_song_chords))
    chords_progression_summary_string += 'Total number of chords: ' + str(len(all_song_chords)) + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'
    print('=' * 70)
    print('Most common pitches chord:', list(Counter(tuple([a[0] for a in all_song_chords])).most_common(1)[0][0]), '===', Counter(tuple([a[0] for a in all_song_chords])).most_common(1)[0][1], 'count')
    chords_progression_summary_string += 'Most common pitches chord: ' + str(list(Counter(tuple([a[0] for a in all_song_chords])).most_common(1)[0][0])) + ' === ' + str(Counter(tuple([a[0] for a in all_song_chords])).most_common(1)[0][1]) + ' count' + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'
    print('=' * 70)
    print('Most common tones chord:', list(Counter(tuple([a[1] for a in all_song_chords])).most_common(1)[0][0]), '===', Counter(tuple([a[1] for a in all_song_chords])).most_common(1)[0][1], 'count')
    chords_progression_summary_string += 'Most common tones chord: ' + str(list(Counter(tuple([a[1] for a in all_song_chords])).most_common(1)[0][0])) + ' === ' + str(Counter(tuple([a[1] for a in all_song_chords])).most_common(1)[0][1]) + ' count' + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'    
    print('=' * 70)
    print('Sorted unique songs chords set:', len(sorted(set(tuple([a[1] for a in all_song_chords])))), 'count')
    chords_progression_summary_string += 'Sorted unique songs chords set: ' + str(len(sorted(set(tuple([a[1] for a in all_song_chords]))))) +  ' count' + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'
    #print('=' * 70)
    for c in sorted(set(tuple([a[1] for a in all_song_chords]))):
        #print(list(c))
        chords_progression_summary_string += str(list(c)) + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'
    print('=' * 70)
    print('Grouped songs chords set:', len(TMIDIX.grouped_set(tuple([a[1] for a in all_song_chords]))), 'count')
    chords_progression_summary_string += 'Grouped songs chords set: ' + str(len(TMIDIX.grouped_set(tuple([a[1] for a in all_song_chords])))) + ' count' + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'
    print('=' * 70)
    for c in TMIDIX.grouped_set(tuple([a[1] for a in all_song_chords])):
        #print(list(c))
        chords_progression_summary_string += str(list(c)) + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'
    #print('=' * 70)
    #print('All songs chords')
    chords_progression_summary_string += 'All songs chords' + '\n'
    chords_progression_summary_string += '=' * 70
    chords_progression_summary_string += '\n'
    #print('=' * 70)
    for i, pc_tc in enumerate(all_song_chords):
        #print('Song chord #', i)
        chords_progression_summary_string += 'Song chord # ' + str(i) + '\n'
        #print(list(pc_tc[0]), '===', list(pc_tc[1]))
        chords_progression_summary_string += str(list(pc_tc[0])) + ' === ' + str(list(pc_tc[1])) + '\n'
        #print('=' * 70)
        chords_progression_summary_string += '=' * 70
        chords_progression_summary_string += '\n'
    #========================================================
    
    print('-' * 70)
    print('Req end time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
    print('-' * 70)
    print('Req execution time:', (reqtime.time() - start_time), 'sec')

    return output_audio, output_plot, output_midi, chords_progression_summary_string

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

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)

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

    soundfont = "SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2"

    print('Loading Monster Harmonized Melodies MIDI Dataset...')
    print('=' * 70)
    all_chords_chunks = TMIDIX.Tegridy_Any_Pickle_File_Reader('Monster_Harmonized_Melodies_MIDI_Dataset')
    
    print('=' * 70)
    print('Total number of harmonized melodies:', len(all_chords_chunks))
    print('=' * 70)
    print('Loading melodies...')

    long_mels_chunks_mult = []
    long_chords_chunks_mult = []
    
    for c in tqdm(all_chords_chunks):
      long_mels_chunks_mult.append([p % 12 for p in c[0]])
      long_chords_chunks_mult.append(c[1])
    
    print('Done!')
    print('=' * 70)
    print('Total loaded melodies count:', len(long_mels_chunks_mult))
    print('=' * 70)
    
    #===============================================================================
    
    app = gr.Blocks()
    with app:
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Monophonic MIDI Melody Harmonizer</h1>")
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Retrieval augmented harmonization of any MIDI melody</h1>")
        gr.Markdown(
            "![Visitors](https://api.visitorbadge.io/api/visitors?path=asigalov61.Monophonic-MIDI-Melody-Harmonizer&style=flat)\n\n"
            "This is a demo for Tegridy MIDI Dataset\n\n"
            "Check out [Tegridy MIDI Dataset](https://github.com/asigalov61/Tegridy-MIDI-Dataset) on GitHub!\n\n"
        )

        gr.Markdown("## Select generation options")

        minimum_song_length_in_chords_chunks = gr.Slider(4, 60, value=30, step=1, label="Minimum song length in chords chunks")
        chords_chunks_memory_ratio = gr.Slider(0, 1, value=1, step=0.1, label="Chords chunks memory ratio")
        chord_time_step = gr.Slider(100, 1000, value=500, step=50, label="Chord time step")
        merge_chords_notes = gr.Slider(0, 4000, value=1000, step=100, label="Merged chords notes max time")
        melody_MIDI_patch_number = gr.Slider(0, 127, value=40, step=1, label="Melody MIDI patch number")
        chords_progression_MIDI_patch_number = gr.Slider(0, 127, value=0, step=1, label="Chords progression MIDI patch number")
        base_MIDI_patch_number = gr.Slider(0, 127, value=35, step=1, label="Base MIDI patch number")

        run_btn = gr.Button("harmonize", variant="primary")

        gr.Markdown("## Generation results")
        
        output_audio = gr.Audio(label="Output MIDI audio", format="mp3", elem_id="midi_audio")
        output_plot = gr.Plot(label="Output MIDI score plot")
        output_midi = gr.File(label="Output MIDI file", file_types=[".mid"])

        output_cp_summary = gr.Textbox(label="Generated chords progression info and stats")

        run_event = run_btn.click(Generate_Chords_Progression, 
                                                              [minimum_song_length_in_chords_chunks,
                                                                chords_chunks_memory_ratio,
                                                                chord_time_step,
                                                                merge_chords_notes,
                                                                melody_MIDI_patch_number,
                                                                chords_progression_MIDI_patch_number,
                                                                base_MIDI_patch_number],                                                                                                                       
                                                               [output_audio, output_plot, output_midi, output_cp_summary]
                                 )
   
        app.queue().launch()