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

from midi_to_colab_audio import midi_to_colab_audio

from huggingface_hub import hf_hub_download

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'

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

def split_by_sublist(main_list, sublist):

    n = len(sublist)
    indices = []
    result = []

    # Find all starting indices where sublist matches main_list
    for i in range(len(main_list) - n + 1):
        if main_list[i:i + n] == sublist:
            indices.append(i)

    if not indices:
        # If the sublist is not found, return the entire main_list as one segment
        return [main_list]

    # Use the indices to split the main_list
    for idx, start_index in enumerate(indices):
        if idx < len(indices) - 1:
            # Slice from current sublist start to the next sublist start
            end_index = indices[idx + 1]
            segment = main_list[start_index:end_index]
        else:
            # For the last sublist occurrence, include the rest of the list
            segment = main_list[start_index:]
        result.append(segment)

    return result

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

def fuzzy_find(values, threshold=256):
    result = set()
    for i in range(len(values)):
        for j in range(i+1, len(values)):
            if abs(values[i] - values[j]) <= threshold:
                result.add(values[j])
                result.add(values[i])
    return sorted(result)

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

def kmp_search(main_list, sublist):
    if not sublist:
        return -1, -1  # Sublist is empty
    
    # Preprocess the sublist to get the longest prefix-suffix array
    lps = [0] * len(sublist)
    length = 0  # Length of the previous longest prefix suffix
    i = 1
    while i < len(sublist):
        if sublist[i] == sublist[length]:
            length += 1
            lps[i] = length
            i += 1
        else:
            if length != 0:
                length = lps[length - 1]
            else:
                lps[i] = 0
                i += 1
    
    # Start the search
    i = j = 0  # Index for main_list and sublist
    while i < len(main_list):
        if main_list[i] == sublist[j]:
            i += 1
            j += 1
        if j == len(sublist):
            start_index = i - j
            end_index = i - 1
            return start_index, end_index
        elif i < len(main_list) and main_list[i] != sublist[j]:
            if j != 0:
                j = lps[j - 1]
            else:
                i += 1
    return -1, -1  # Sublist not found

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

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 Remix_MIDI(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(Remix_MIDI, 
                   [input_midi,
                    chords_chunks_len,
                    num_mix_chunks
                   ], 
                   [output_audio,
                    output_plot,
                    output_midi                          
                   ]
                  )
 
#==================================================================================

demo.launch()

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