# =================================================================================================
# https://huggingface.co/spaces/asigalov61/Chords-Progressions-Generator
# =================================================================================================

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

import gradio as gr

import numpy as np

import os
import random
from collections import Counter
import math
from tqdm import tqdm

import TMIDIX

from midi_to_colab_audio import midi_to_colab_audio

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

def Generate_Chords_Progression(total_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('Total song length in chords chunks:', total_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(len(matching_long_chords_chunks))
      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 processed Pitches Chords Progressions dataset data...')
    print('=' * 70)
    long_tones_chords_dict, all_long_chords_tokens_chunks, all_long_good_chords_chunks = TMIDIX.Tegridy_Any_Pickle_File_Reader('processed_chords_progressions_chunks_data')
    
    print('=' * 70)
    print('Resulting chords dictionary size:', len(long_tones_chords_dict))
    print('=' * 70)
    print('Loading chords chunks...')

    chunk_size = 4
    
    src_long_chunks = np.array([a[:chunk_size] for a in all_long_chords_tokens_chunks])
    
    print('Done!')
    print('=' * 70)
    print('Total chords chunks count:', len(all_long_good_chords_chunks))
    print('=' * 70)
    
    #===============================================================================
    
    app = gr.Blocks()
    with app:
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Chords Progressions Generator</h1>")
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Generate unique chords progressions</h1>")
        gr.Markdown(
            "![Visitors](https://api.visitorbadge.io/api/visitors?path=asigalov61.Chords-Progressions-Generator&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"
            "[Open In Colab]"
            "(https://colab.research.google.com/github/asigalov61/Tegridy-MIDI-Dataset/blob/master/Chords-Progressions/Pitches_Chords_Progressions_Generator.ipynb)"
            " for all options, faster execution and endless generation"
        )

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

        total_song_length_in_chords_chunks = gr.Slider(4, 60, value=30, step=1, label="Total 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("generate", 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, 
                                                              [total_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()