#==================================================================================
# https://huggingface.co/spaces/asigalov61/MIDI-Genre-Classifier
#==================================================================================

print('=' * 70)
print('MIDI Genre Classifier Gradio App')

print('=' * 70)
print('Loading core MIDI Genre Classifier modules...')

import os
import copy

import time as reqtime
import datetime
from pytz import timezone

print('=' * 70)
print('Loading main MIDI Genre Classifier modules...')

os.environ['USE_FLASH_ATTENTION'] = '1'

import torch

torch.set_float32_matmul_precision('medium')
torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(True)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_cudnn_sdp(True)

import numpy as np

from huggingface_hub import hf_hub_download

import TMIDIX

from midi_to_colab_audio import midi_to_colab_audio

from x_transformer_1_23_2 import *

import random

import tqdm

print('=' * 70)
print('Loading aux MIDI Genre Classifier modules...')

import matplotlib.pyplot as plt

import gradio as gr
import spaces

print('=' * 70)
print('PyTorch version:', torch.__version__)
print('=' * 70)
print('Done!')
print('Enjoy! :)')
print('=' * 70)

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

MODEL_CHECKPOINT = 'Giant_Music_Transformer_Medium_Trained_Model_42174_steps_0.5211_loss_0.8542_acc.pth'

SOUDFONT_PATH = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'

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

print('=' * 70)
print('Loading MIDI GAS processed scores dataset...')

midi_gas_ps_pickle = hf_hub_download(repo_id='asigalov61/MIDI-GAS', 
                                     filename='MIDI_GAS_Processed_Scores_CC_BY_NC_SA.pickle',
                                     repo_type='dataset'
                                    )

midi_gas_ps = TMIDIX.Tegridy_Any_Pickle_File_Reader(midi_gas_ps_pickle)

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

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

print('=' * 70)
print('Loading MIDI GAS processed scores embeddings dataset...')

midi_gas_pse_pickle = hf_hub_download(repo_id='asigalov61/MIDI-GAS', 
                                      filename='MIDI_GAS_Processed_Scores_Embeddings_CC_BY_NC_SA.pickle',
                                      repo_type='dataset'
                                     )

midi_gas_pse = np.array([a[3] for a in TMIDIX.Tegridy_Any_Pickle_File_Reader(midi_gas_pse_pickle)])

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

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

print('=' * 70)
print('Instantiating model...')

device_type = 'cuda'
dtype = 'bfloat16'

ptdtype = {'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)

SEQ_LEN = 8192
PAD_IDX = 19463

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

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

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

model_checkpoint = hf_hub_download(repo_id='asigalov61/Giant-Music-Transformer', filename=MODEL_CHECKPOINT)

model.load_state_dict(torch.load(model_checkpoint, map_location='cpu', weights_only=True))

print('=' * 70)
print('Done!')
print('=' * 70)
print('Model will use', dtype, 'precision...')
print('=' * 70)

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

def load_midi(input_midi):

    raw_score = TMIDIX.midi2single_track_ms_score(input_midi)
    
    escore = TMIDIX.advanced_score_processor(raw_score, return_enhanced_score_notes=True)[0]
    
    escore_notes = TMIDIX.augment_enhanced_score_notes(escore)
    
    instruments_list = list(set([y[6] for y in escore_notes]))
    
    tok_score = []
    
    if 128 in instruments_list:
        drums_present = 19331
    
    else:
        drums_present = 19330
    
    pat = escore_notes[0][6]
    
    tok_score.extend([19461, drums_present, 19332+pat])
    
    tok_score.extend(TMIDIX.multi_instrumental_escore_notes_tokenized(escore_notes)[:8190])

    return tok_score

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

def logsumexp_pooling(x, dim=1, keepdim=False):
    
    max_val, _ = torch.max(x, dim=dim, keepdim=True)
    
    lse = max_val + torch.log(torch.mean(torch.exp(x - max_val), dim=dim, keepdim=keepdim) + 1e-10)
    
    return lse

def gem_pooling(x, p=3.0, eps=1e-6):
    
    pooled = torch.mean(x ** p, dim=1)
    
    return pooled.clamp(min=eps).pow(1 / p)

def median_pooling(x, dim=1):
    return torch.median(x, dim=dim).values

def rms_pooling(x, dim=1):
    return torch.sqrt(torch.mean(x ** 2, dim=dim) + 1e-6)

def get_embeddings(inputs):
    
    with ctx:
        with torch.no_grad():
            out = model(inputs)
    
    cache = out[2]

    hidden = cache.layer_hiddens[-1]
    
    mean_pool   = torch.mean(hidden, dim=1)
    max_pool    = torch.max(hidden, dim=1).values
    lse_pool    = logsumexp_pooling(hidden, dim=1)
    gem_pool    = gem_pooling(hidden, p=3.0)
    median_pool = median_pooling(hidden, dim=1)
    rms_pool    = rms_pooling(hidden, dim=1)
    
    concat_pool = torch.cat((mean_pool,
                             max_pool,
                             lse_pool[0][:, :512],
                             gem_pool[:, :512],
                             median_pool[:, :512],
                             rms_pool[:, :512]), dim=1)
    
    return concat_pool.cpu().detach().numpy()[0]

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

def cosine_similarity_numpy(src_array, trg_array):
   
    src_norm = np.linalg.norm(src_array)

    trg_norms = np.linalg.norm(trg_array, axis=1)

    dot_products = np.dot(trg_array, src_array)

    cosine_sims = dot_products / (src_norm * trg_norms + 1e-10)
    
    return cosine_sims.tolist()

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

def select_best_output(outputs, embeddings, src_embeddings, top_k=10):

    emb_sims = cosine_similarity_numpy(src_embeddings, embeddings)

    sorted_emb_sims = sorted(emb_sims, reverse=True)

    hits = []
    hits_idxs = []
    
    for s in sorted_emb_sims[:top_k]:
        idx = emb_sims.index(s)
        hits_idxs.append(idx)
        hits.extend([[str(s)] + outputs[idx][:3]])
        
    return hits, hits_idxs

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

@spaces.GPU
def Classify_MIDI_Genre(input_midi):

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

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

    print('=' * 70)
    print('Requested settings:')
    print('=' * 70)

    fn = os.path.basename(input_midi)
    fn1 = fn.split('.')[0]
    
    print('Input MIDI file name:', fn)
    print('=' * 70)

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

    model.to(device_type)
    model.eval()

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

    print('Loading and prepping source MIDI...')

    src_score = load_midi(input_midi.name)

    inp = torch.LongTensor([src_score]).to(device_type)

    src_emb = get_embeddings(inp)

    print('Done!')

    #===============================================================================
        
    print('Sample embeddings values', src_emb[:3])

    #===============================================================================
    
    print('=' * 70)
    print('Classifying...')

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

    result = select_best_output(midi_gas_ps, midi_gas_pse, src_emb)

    results_str = ''

    for i, r in enumerate(result[0]):
        print(' --- '.join([str(i+1).zfill(2)] + r))
        results_str += ' --- '.join([str(i+1).zfill(2)] + r) + '\n'

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

    song_name = ' --- '.join(midi_gas_ps[result[1][0]][:3])
    print('Song entry', song_name)
    
    song = midi_gas_ps[result[1][0]][3]
    print('Sample INTs', song[:15])
    print('=' * 70)

    song_f = []
    
    if len(song) != 0:
    
        time = 0
        dur = 0
        vel = 90
        pitch = 0
        channel = 0
        
        patches = [-1] * 16
        patches[9] = 9
        
        channels = [0] * 16
        channels[9] = 1
        
        for ss in song:
        
          if 0 <= ss < 256:
        
              time += ss * 16
        
          if 256 <= ss < 2304:
        
              dur = ((ss-256) // 8) * 16
              vel = (((ss-256) % 8)+1) * 15
        
          if 2304 <= ss < 18945:
        
              patch = (ss-2304) // 129
        
              if patch < 128:
        
                  if patch not in patches:
                    if 0 in channels:
                        cha = channels.index(0)
                        channels[cha] = 1
                    else:
                        cha = 15
        
                    patches[cha] = patch
                    channel = patches.index(patch)
                  else:
                    channel = patches.index(patch)
        
              if patch == 128:
                  channel = 9
        
              pitch = (ss-2304) % 129
        
              song_f.append(['note', time, dur, channel, pitch, vel, patch ])
        
        patches = [0 if x==-1 else x for x in patches]


    fn1 = "MIDI-Genre-Classifier-Composition"
    
    detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(song_f,
                                                              output_signature = 'MIDI Genre Classifier',
                                                              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=SOUDFONT_PATH,
                                sample_rate=16000,
                                volume_scale=10,
                                output_for_gradio=True
                                )
    
    print('Done!')
    print('=' * 70)

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

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

    output_cls_results = str(results_str)

    print('Output MIDI file name:', output_midi)
    print('Output MIDI melody title:', output_title)
    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_title, output_audio, output_plot, output_midi, output_cls_results
    
#==================================================================================

PDT = timezone('US/Pacific')

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

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

with gr.Blocks() as demo:

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

    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>MIDI Genre Classifier</h1>")
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Accurately classify any MIDI by top music genre</h1>")
    
    gr.HTML("""            
            <p> 
                <a href="https://huggingface.co/spaces/asigalov61/MIDI-Genre-Classifier?duplicate=true">
                    <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-md.svg" alt="Duplicate in Hugging Face">
                </a>
            </p>
            
            for faster execution and endless classification!
            """)
    
    #==================================================================================
    
    gr.Markdown("## Upload any MIDI or select an example MIDI below")
    
    input_midi = gr.File(label="Input MIDI", 
                         file_types=[".midi", ".mid", ".kar"]
                        )
    
    generate_btn = gr.Button("Classify", variant="primary")

    gr.Markdown("## Classification results")

    output_title = gr.Textbox(label="MIDI title")
    output_audio = gr.Audio(label="MIDI audio", format="wav", elem_id="midi_audio")
    output_plot = gr.Plot(label="MIDI score plot")
    output_midi = gr.File(label="MIDI file", file_types=[".mid"])
    output_cls_results = gr.Textbox(label="MIDI classification results")

    generate_btn.click(Classify_MIDI_Genre, 
                       [input_midi
                       ], 
                       [output_title,
                        output_audio,
                        output_plot,
                        output_midi,
                        output_cls_results
                       ]
                      )
    
    gr.Examples(
                [["Hotel California.mid"],
                 ["Come To My Window.mid"]
                ],
                [input_midi
                ],
                [output_title,
                 output_audio,
                 output_plot,
                 output_midi,
                 output_cls_results
                ],
                Classify_MIDI_Genre
    )
              
#==================================================================================

demo.launch()

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