#==================================================================
# https://huggingface.co/spaces/asigalov61/Popular-Hook-Transformer
#==================================================================

import time as reqtime
import datetime
from pytz import timezone

import statistics
import re
import tqdm

import gradio as gr
import spaces

from x_transformer_1_23_2 import *
import random

from midi_to_colab_audio import midi_to_colab_audio
import TMIDIX

import matplotlib.pyplot as plt
    
#=====================================================================================

print('=' * 70)
print('Popular Hook Transformer')
print('=' * 70)

print('Loading Popular Hook Transformer training data...')
print('=' * 70)

melody_chords_f = TMIDIX.Tegridy_Any_Pickle_File_Reader('Popular_Hook_Transformer_Training_Data.pickle')

print('=' * 70)

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

SEQ_LEN = 512
PAD_IDX = 918
DEVICE = 'cuda'

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

def str_strip(string):
    return re.sub(r'[^A-Za-z-]+', '', string).rstrip('-')

def mode_time(seq):
    return statistics.mode([t for t in seq if 0 < t < 128])

def mode_dur(seq):
    return statistics.mode([t-128 for t in seq if 128 < t < 256])

def mode_pitch(seq):
    return statistics.mode([t % 128 for t in seq if 256 < t < 512])

sections_dict = sorted(set([str_strip(s[2]).rstrip('-') for s in melody_chords_f]))

train_data = []

for m in tqdm.tqdm(melody_chords_f):
    
    if 64 < len(m[5]) < 506:
        
        for tv in range(-3, 3):
        
            section = str_strip(m[2])
            section_tok = sections_dict.index(section)
            
            score = [t+tv if 256 < t < 512 else t for t in m[5]]

            seq = [916] + [section_tok+512, mode_time(score)+532, mode_dur(score)+660, mode_pitch(score)+tv+788]

            seq += score

            seq += [917]

            seq = seq + [PAD_IDX] * (SEQ_LEN - len(seq))

            train_data.append(seq)

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

print('Done!')
print('=' * 70)
print('All data is good:', len(max(train_data, key=len)) == len(min(train_data, key=len)))
print('=' * 70)
print('Randomizing training data...')
random.shuffle(train_data)
print('Done!')
print('=' * 70)
print('Total length of training data:', len(train_data))
print('=' * 70)

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

print('Loading Popular Hook Transformer pre-trained model...')
print('=' * 70)

print('Instantiating model...')

model = TransformerWrapper(
    num_tokens = PAD_IDX+1,
    max_seq_len = SEQ_LEN,
    attn_layers = Decoder(dim = 1024, 
                          depth = 4, 
                          heads = 32,
                          rotary_pos_emb = True,
                          attn_flash = True
                         )
    )

model = AutoregressiveWrapper(model, ignore_index = PAD_IDX, pad_value=PAD_IDX)

print('=' * 70)
print('Loading model checkpoint...')

model_path = 'Popular_Hook_Transformer_Small_Trained_Model_10869_steps_0.2308_loss_0.9252_acc.pth'

model.load_state_dict(torch.load(model_path, map_location='cpu'))

print('Done!')
print('=' * 70)

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

@spaces.GPU
def Generate_POP_Section(input_comp_section, 
                         input_mode_time, 
                         input_mode_dur,
                         input_mode_ptc
                         ):
    
    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('Composition section:', input_comp_section)
    print('Mode time:', input_mode_time)
    print('Mode duration:', input_mode_dur)
    print('Mode pitch:', input_mode_ptc)
    print('=' * 70)

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

    print('Generating...')
    
    if input_comp_section == 'random':
        seq = [916]

    else:
        seq = [916, sections_dict.index(input_comp_section)+512]

        input_seq = [input_mode_time, input_mode_dur, input_mode_ptc]
        input_seq_toks = [input_mode_time+532, input_mode_dur+660, input_mode_ptc+788]

        if 0 in input_seq:
            input_seq = input_seq_toks[:input_seq.index(0)]

        else:
            input_seq = input_seq_toks

        seq += input_seq

    model.to(DEVICE)
    model.eval()

    x = torch.LongTensor(seq).to(DEVICE)
    
    with torch.amp.autocast(device_type=DEVICE, dtype=torch.bfloat16):
        
        out = model.generate(x,
                             512-len(seq),
                             temperature=0.9,
                             filter_logits_fn=top_p,
                             filter_kwargs={'thres': 0.96},
                             eos_token=917,
                             return_prime=True,
                             verbose=True)

    song = out.tolist()[0]

    print('Done!')
    print('=' * 70)
        
    #===============================================================================
    
    print('Rendering results...')
    
    print('=' * 70)
    
    comp_section = sections_dict[song[1]-512]
    comp_mode_time = song[2]-532
    comp_mode_dur = song[3]-660
    comp_mode_ptc = song[4]-788
    
    comp_summary = ''

    comp_summary += 'Generated section: ' + str(comp_section) + '\n'
    comp_summary += 'Generated mode time: ' + str(comp_mode_time) + '\n'
    comp_summary += 'Generated mode duration: ' + str(comp_mode_dur) + '\n'
    comp_summary += 'Generated mode pitch :' + str(comp_mode_ptc)
    
    print('Sample INTs', song[:5])
    print('=' * 70)

    song_f = []

    time = 0
    dur = 0
    vel = 90
    pitch = 0
    channel = 0
    
    for ss in song:
    
      if 0 <= ss < 128:
    
          time += ss * 32
    
      if 128 <= ss < 256:
    
          dur = (ss-128)* 32
    
      if 256 <= ss < 512:

          pitch = (ss-256) % 128
          cha = (ss-256) // 128

          if cha == 0:
              channel = 3
              vel = 110+(pitch % 12)
              patch = 40
              
          else:
              channel = 0
              vel = max(40, pitch)
              patch = 0
    
          song_f.append(['note', time, dur, channel, pitch, vel, patch ])

    fn1 = 'Popular-Hook-Transformer-Composition'

    detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
                                                              output_signature = 'Popular Hook Transformer',
                                                              output_file_name = fn1,
                                                              track_name='Project Los Angeles'
                                                              )
    
    new_fn = fn1+'.mid'
            
    
    audio = midi_to_colab_audio(new_fn, 
                        soundfont_path=soundfont,
                        sample_rate=16000,
                        volume_scale=10,
                        output_for_gradio=True
                        )
    
    print('Done!')
    print('=' * 70)

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

    output_midi_title = str(fn1)
    output_midi = str(new_fn)
    output_audio = (16000, audio)
    
    output_plot = TMIDIX.plot_ms_SONG(song_f, plot_title=output_midi_title, return_plt=True)

    print('Output MIDI file name:', output_midi)
    print('Output MIDI title:', output_midi_title)
    print('Output MIDI summary', comp_summary)
    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_title, comp_summary, output_midi, output_audio, output_plot

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

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"

    app = gr.Blocks()
    
    with app:
        
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Popular Hook Transformer</h1>")
        gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Generate unique POP music sections</h1>")
        gr.Markdown(
            "This is a demo for popular-hook MIDI Dataset\n\n"
            "Check out [popular-hook](https://huggingface.co/datasets/NEXTLab-ZJU/popular-hook) on Hugging Face!\n\n"
        )
        
        gr.Markdown("## Select POP composition section to generate:")

        input_comp_section = gr.Dropdown(sections_dict + ['random'], label="Composition section", value='random')
        
        gr.Markdown("## Select generation options:")

        input_mode_time = gr.Slider(0, 127, value=0, step=1, label="Composition mode time")
        input_mode_dur = gr.Slider(0, 127, value=0, step=1, label="Composition mode dur")
        input_mode_ptc = gr.Slider(0, 127, value=0, step=1, label="Composition mode pitch")
        
        run_btn = gr.Button("Generate", variant="primary")
        
        gr.Markdown("## Output results")

        output_midi_title = gr.Textbox(label="Output MIDI title")
        output_midi_summary = gr.Textbox(label="Output MIDI summary")
        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"])

        run_event = run_btn.click(Generate_POP_Section, [input_comp_section, 
                                                        input_mode_time, 
                                                        input_mode_dur,
                                                        input_mode_ptc
                                                        ],
                                                        [output_midi_title,
                                                         output_midi_summary,
                                                         output_midi, 
                                                         output_audio, 
                                                         output_plot]
                                                        )

        gr.Examples([["intro", 10, 15, 60],
                    ],
                    [input_comp_section, 
                    input_mode_time, 
                    input_mode_dur,
                    input_mode_ptc
                    ],
                    [output_midi_title,
                    output_midi_summary,
                    output_midi, 
                    output_audio, 
                    output_plot],
                    Generate_POP_Section,
                    cache_examples=True,
                    cache_mode='eager'
                    )
        
        app.queue().launch()