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

print('=' * 70)
print('MIDI Remixer Gradio App')

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

import os
import copy
import statistics
import random
from collections import Counter

import time as reqtime
import datetime
from pytz import timezone

import tqdm

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

import TMIDIX

import numpy as np

from midi_to_colab_audio import midi_to_colab_audio

from huggingface_hub import hf_hub_download
from datasets import load_dataset

import gradio as gr

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

import matplotlib.pyplot as plt

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

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

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

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

print('=' * 70)

midi_cores_dataset = load_dataset("asigalov61/MIDI-Cores")

print('=' * 70)
print('Prepping MIDI cores data...')
print('=' * 70)

all_core_chords = []

for entry in tqdm.tqdm(midi_cores_dataset['train']):
    all_core_chords.append(entry['core_chords'])

all_core_chords = np.array(all_core_chords)

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

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

def load_midi(midi_file):

    print('Loading MIDI...')

    raw_score = TMIDIX.midi2single_track_ms_score(midi_file)
    escore_notes = TMIDIX.advanced_score_processor(raw_score, return_enhanced_score_notes=True)[0]
    escore_notes = TMIDIX.augment_enhanced_score_notes(escore_notes, sort_drums_last=True)
    
    zscore = TMIDIX.recalculate_score_timings(escore_notes)
    
    cscore = TMIDIX.chordify_score([1000, zscore])
    
    tones_chords = []
    tones_chords_idxs = []
    
    for i, c in enumerate(cscore):
        pitches = [e[4] for e in c if e[3] != 9]
    
        if pitches:
            tones_chord = sorted(set([p % 12 for p in pitches]))
                 
            if tones_chord in TMIDIX.ALL_CHORDS_SORTED:
                chord_token = TMIDIX.ALL_CHORDS_SORTED.index(tones_chord)
                
            else:
                tones_chord = TMIDIX.check_and_fix_tones_chord(tones_chord)
                chord_token = TMIDIX.ALL_CHORDS_SORTED.index(tones_chord)
        
            tones_chords.append(chord_token)
            tones_chords_idxs.append(i)
    
    print('Done!')
    print('=' * 70)
    print('MIDI has', len(tones_chords), 'chords')
    print('=' * 70)

    return cscore, tones_chords, tones_chords_idxs

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

def find_max_exact_match(src, trg):
    """
    Find the row in the 2D trg array that has the maximum number of exact element matches with the 1D src array.

    Parameters:
    src (numpy.ndarray): 1D source array.
    trg (numpy.ndarray): 2D target array.

    Returns:
    int: Index of the row in trg that has the maximum number of exact matches with src.
    """
    # Compare src with each row in trg and count exact matches
    match_counts = np.sum(trg == src, axis=1)
    
    # Find the index of the row with the maximum number of matches
    max_match_idx = np.argmax(match_counts)
    
    return max_match_idx

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

def Match_Cores(input_midi,
                chords_chunks_len,
                num_mix_chunks
               ):

    #===============================================================================
    
    print('=' * 70)
    print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))
    start_time = reqtime.time()
    print('=' * 70)
    
    fn = os.path.basename(input_midi)
    fn1 = fn.split('.')[0]

    print('=' * 70)
    print('Requested settings:')
    print('=' * 70)
    print('Input MIDI file name:', fn)
    print('Chords chunks len:', chords_chunks_len)
    print('Num mix chunks:', num_mix_chunks)
   
    print('=' * 70)

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

    cscore, tones_chords, tones_chords_idxs = load_midi(input_midi.name)
    
    #===============================================================================
    
    print('Prepping chords chunks...')

    chunk_len = chords_chunks_len
    
    all_chords_chunks = Counter()
    
    for i in range(0, len(tones_chords)):
        chunk = tones_chords[i:i+chunk_len]
        if len(chunk) == chunk_len:
            all_chords_chunks[tuple(chunk)] += 1

    print('Done!')
    print('=' * 70)
    print('Number of chords chunks:', len(all_chords_chunks))
    print('=' * 70)

    #==================================================================
    
    print('Remixing MIDI...')
    print('=' * 70)

    other_chunks = [list(k) for k, v in all_chords_chunks.most_common(10)]

    all_match_chunks = {}

    for mc_chunk in other_chunks:
        chunk_lens = fuzzy_find([len(l) for l in split_by_sublist(tones_chords, mc_chunk) if len(l) > chunk_len])
        match_chunks = [l for l in split_by_sublist(tones_chords, mc_chunk) if len(l) in chunk_lens]
        all_match_chunks[tuple(mc_chunk)] = match_chunks

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

    chunks_map = random.choices(other_chunks, k=num_mix_chunks)

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

    matched_chunks = []
    start_chunk = []
    res = -1
    tries = 0
    while  res == -1 and len(start_chunk[:res]) < chunk_len and tries < 100:
    
        start_chunk = random.choice(all_match_chunks[tuple(chunks_map[0])])
        res = kmp_search(start_chunk, chunks_map[1])[0]
    
        tries += 1
    
        #print(res)
    if tries < 100:
        matched_chunks.append(start_chunk[:res])
    
    else:
        print('FAIL!')

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

    mchunks = [-1]
    tries = 0
    while -1 in mchunks and tries < 100:
    
        mchunks = []
        
        for i, chunk in enumerate(chunks_map[1:-1]):
            start_chunk = random.choice(all_match_chunks[tuple(chunk)])
            res = kmp_search(start_chunk, chunks_map[i+2])[0]
        
            #print(res)
            mchunks.append(start_chunk[:res])
    
        tries += 1
    
    if tries < 100:
        matched_chunks.extend(mchunks)
    
    else:
        print('FAIL!')

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

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

    print('Creating final MIDI score...')

    all_matches_dscores = []

    for match in matched_chunks:
    
        start, end = kmp_search(tones_chords, match)
        
        sidx = tones_chords_idxs[start]
        eidx = tones_chords_idxs[end]
    
        all_matches_dscores.append(TMIDIX.delta_score_notes(TMIDIX.flatten(cscore[sidx:eidx+1])))

    #===============================================================================
    
    new_dscore = all_matches_dscores[0]
    
    for score in all_matches_dscores[1:]:
        score[0][1] = statistics.mode([e[1] for e in new_dscore[-75:] if e[1] != 0 and e[3] != 9])
        new_dscore.extend(score)
    
    new_escore = TMIDIX.delta_score_to_abs_score(new_dscore)
        
    #===============================================================================

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

    #===============================================================================    
    
    print('Rendering results...')
    
    print('=' * 70)
    print('Sample MIDI events:', new_escore[:3])
    print('=' * 70)

    output_score, patches, overflow_patches = TMIDIX.patch_enhanced_score_notes(new_escore)

    fn1 = "MIDI-Remixer-Composition"
    
    detailed_stats = TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(output_score,
                                                              output_signature = 'MIDI Remixer',
                                                              output_file_name = fn1,
                                                              track_name='Project Los Angeles',
                                                              list_of_MIDI_patches=patches,
                                                              timings_multiplier=16
                                                              )
    
    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_midi = str(new_fn)
    output_audio = (16000, audio)
    
    output_plot = TMIDIX.plot_ms_SONG(output_score, 
                                      plot_title=output_midi,
                                      timings_multiplier=16,
                                      return_plt=True
                                     )

    print('Output MIDI file name:', output_midi)
    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_audio, output_plot, output_midi
    
#==================================================================================

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 Remixer</h1>")
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>Remix repeating parts of any MIDI into one song</h1>")
    gr.HTML("""            
            <p> 
                <a href="https://huggingface.co/spaces/asigalov61/MIDI-Remixer?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("## Upload source MIDI")
    
    input_midi = gr.File(label="Input MIDI", file_types=[".midi", ".mid", ".kar"])
    
    gr.Markdown("## Mixing options")
    
    chords_chunks_len = gr.Slider(3, 7, value=4, step=1, label="Number of chords to match for each repeating chunk")
    num_mix_chunks = gr.Slider(4, 20, value=10, step=1, label="Number of repeating chunks to mix")
    
    mix_btn = gr.Button("Remix", variant="primary")

    gr.Markdown("## Mixing results")

    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"])

    mix_btn.click(Match_Cores, 
                   [input_midi,
                    chords_chunks_len,
                    num_mix_chunks
                   ], 
                   [output_audio,
                    output_plot,
                    output_midi                          
                   ]
                  )
 
#==================================================================================

demo.launch()

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