# =================================================================================================
# 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

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

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,
                                add_drums,
                                output_as_solo_piano
                               ):

    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('-' * 70)
    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('Add drum track:', add_drums)
    print('Output as solo Piano:', output_as_solo_piano)
    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)

    if add_drums and not output_as_solo_piano:
        output_score = TMIDIX.augment_enhanced_score_notes(output_score)
        output_score = TMIDIX.advanced_add_drums_to_escore_notes(output_score)

        for e in output_score:
            e[1] *= 16
            e[2] *= 16
            
        ctimes = sorted(set([e[1] for e in output_score if e[3] != 9]))

        for i, e in enumerate(chords_labels[:len(ctimes)]):
            e[1] = ctimes[i]
            
    if output_as_solo_piano:
        output_score = TMIDIX.solo_piano_escore_notes(output_score)

    midi_score = sorted(chords_labels + output_score, key=lambda x: x[1])

    fn1 = "MuseCraft-Chords-Progression-Composition"
    
    detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(midi_score,
                                                              output_signature = 'MuseCraft Chords Progression',
                                                              output_file_name = fn1,
                                                              track_name='Project Los Angeles',
                                                              list_of_MIDI_patches=patches
                                                              )
    
    new_fn = fn1+'.mid'
                
    #========================================================

    output_midi = str(new_fn)
    
    print('Done!')
    print('=' * 70)
    
    #========================================================
   
    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_midi

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

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 = 6
    
    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: left; margin-bottom: 1rem'>Chords Progressions Generator</h1>")
        gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>Generate unique chords progressions</h1>")
        gr.HTML("""
            Check out <a href="https://github.com/MIDIAI/MuseCraft">MuseCraft project</a> on GitHub
            <p>
                <a href="https://huggingface.co/spaces/projectlosangeles/MuseCraft-Chords-Progressions?duplicate=true">
                    <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-md.svg" alt="Duplicate in Hugging Face">
                </a>
            </p>
        """)

        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=250, step=50, label="Chord time step")
        merge_chords_notes = gr.Slider(0, 4000, value=2000, 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")
        add_drums = gr.Checkbox(value=True, label="Add drum track")
        output_as_solo_piano = gr.Checkbox(value=False, label="Output as solo Piano")

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

        gr.Markdown("## Generation results")
        
        output_midi = gr.File(label="Output MIDI file", file_types=[".mid"])

        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,
                                                               add_drums,
                                                               output_as_solo_piano
                                                              ],                                                                                                                       
                                                               [output_midi]
                                 )
   
        app.queue().launch()