#====================================================================================
# https://huggingface.co/spaces/projectlosangeles/Masked-Chordified-Piano-Transformer
#====================================================================================

"""
Masked Chordified Piano Transformer Gradio App
"""

#===================================================================
# Environment requirements (fully cross platform and minimal)
#===================================================================
# pip requirements 
#-------------------------------------------------------------------
# !pip install tqdm
# !pip install numpy
# !pip install matplotlib
# !pip install gradio
# !pip install hf-transfer
# !pip install huggingface_hub
# !pip install torch
# !pip install einops
# !pip install einx
# !pip install scikit-learn
#===================================================================
# apt requirements
#-------------------------------------------------------------------
# !sudo apt install fluidsynth -y
#===================================================================
# Required modules
#-------------------------------------------------------------------
# Download modules from https://github.com/asigalov61/tegridy-tools
#-------------------------------------------------------------------
# TMIDIX.py
# x_transformer_2_3_1.py
# midi_to_colab_audio.py
#===================================================================

# -----------------------------
# CONFIGURATION & GLOBALS
# -----------------------------
TIME_ZONE = 'US/Pacific'

SEQ_LEN = 4096
PAD_IDX = 717

MODELS_CHECKPOINTS = [
    {
        'checkpoint_tag': 'Base Model',
        'checkpoint_name': 'Chordified_Piano_Transformer_Trained_Model_16852_steps_0.3286_loss_0.9059_acc.pth',
        'checkpoint_depth': 12,
        'checkpoint_heads': 16
    },
]

MODEL_DEVICE = 'cuda'

SOUNDFONT_BANK = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'
AUDIO_SAMPLE_RATE = 16000
AUDIO_FORMAT = 'mp3'

NUM_OUT_BATCHES = 10
PREVIEW_LENGTH = 120  # in tokens

OUTPUT_MIDIS_DIR = 'output_midis'

# -----------------------------
# START-UP INFO FUNCTIONS
# -----------------------------
SEP = '=' * 70

def print_sep():
    print(SEP)

print_sep()
print("Masked Chordified Piano Transformer Gradio App")
print_sep()
print("Loading modules...")

# -----------------------------
# ENVIRONMENT & MODULES IMPORTS
# -----------------------------

import os

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

RUNNING_IN_SPACE = (
    os.environ.get("SYSTEM", "").lower() == "spaces"
    or "SPACE_ID" in os.environ
    or "HF_SPACE_ID" in os.environ
)

import argparse

from pathlib import Path
from io import BytesIO

import time as reqtime
import datetime
from pytz import timezone

PDT = timezone(TIME_ZONE)

import random

if RUNNING_IN_SPACE:
    import spaces
    GPU = spaces.GPU
else:
    def GPU(*args, **kwargs):
        def wrapper(fn):
            return fn
        return wrapper
        
import gradio as gr

import TMIDIX

from midi_to_colab_audio import midi_to_colab_audio

import matplotlib.pyplot as plt

from huggingface_hub import hf_hub_download

# -----------------------------
# PyTorch
# -----------------------------

import torch

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

torch.set_float32_matmul_precision('high')
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
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)

MODEL_DTYPE = torch.bfloat16

# -----------------------------
# X-Transformer
# -----------------------------

from x_transformer_2_3_1 import TransformerWrapper, AutoregressiveWrapper, Decoder, top_p

print_sep()
print("PyTorch version:", torch.__version__)
print("Done loading modules!")
print_sep()

# -----------------------------
# SPACES AND LOCAL ARGS
# -----------------------------

def parse_local_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--soundfont-name", type=str, default="SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2")
    return parser.parse_args()

args = parse_local_args() if not RUNNING_IN_SPACE else None

if args:
    SOUNDFONT_BANK = args.soundfont_name

# -----------------------------
# MODELS INIT FUNCTIONS
# -----------------------------
print_sep()

#------------------------------------------------------------------------

def load_model(model_dic):

    print('Instantiating model...')
    
    model = TransformerWrapper(
        num_tokens=PAD_IDX + 1,
        max_seq_len=SEQ_LEN,
        attn_layers=Decoder(
            dim=2048,
            depth=model_dic['checkpoint_depth'],
            heads=model_dic['checkpoint_heads'],
            rotary_pos_emb=True,
            attn_flash=True
        )
    )
    model = AutoregressiveWrapper(model,
                                  ignore_index=PAD_IDX,
                                  pad_value=PAD_IDX
                                 )
    
    print('Done!')
    print_sep()
    print("Model will use", MODEL_DTYPE.__repr__().split('.')[-1], "precision...")
    print("Model will use", MODEL_DEVICE, "device...")
    print_sep()
    print("Loading model checkpoint...")
    print('Checkpoint name:', model_dic['checkpoint_name'])
    print_sep()
    
    checkpoint = hf_hub_download(
        repo_id='asigalov61/Chordified-Piano-Transformer',
        filename=model_dic['checkpoint_name']
    )
    
    model.load_state_dict(torch.load(checkpoint, map_location='cpu'))
    
    model.eval()
    
    model.cpu()
    
    model = torch.compile(model, mode='max-autotune')
    
    print_sep()
    print("Done!")
    print_sep()

    return model_dic['checkpoint_tag'], model

#------------------------------------------------------------------------

models_dict = {}

for model_dic in MODELS_CHECKPOINTS:
    tag, model = load_model(model_dic)
    models_dict[tag] = model

#------------------------------------------------------------------------

ctx = torch.amp.autocast(device_type=MODEL_DEVICE,
                         dtype=MODEL_DTYPE
                        )

print_sep()
print("Done!")
print_sep()

# -----------------------------
# SOUNDFONT LOADING FUNCTION
# -----------------------------
print('Loading SoundFont...')
print_sep()

SOUNDFONT_PATH = hf_hub_download(repo_id='projectlosangeles/soundfonts4u',
                                 repo_type='dataset',
                                 filename=SOUNDFONT_BANK
                                )

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

# -----------------------------
# MIDI PROCESSING FUNCTIONS
# -----------------------------
def load_midi(input_midi):
    
    """Process the input MIDI file and create a token sequence."""

    raw_score = TMIDIX.midi2single_track_ms_score(input_midi.name)
    
    escore_notes = TMIDIX.advanced_score_processor(raw_score, 
                                                   return_enhanced_score_notes=True, 
                                                   apply_sustain=True
                                                  )

    if escore_notes and escore_notes[0]:
    
        sp_escore_notes = TMIDIX.solo_piano_escore_notes(escore_notes[0])
        zscore = TMIDIX.recalculate_score_timings(sp_escore_notes)
        
        escore = TMIDIX.augment_enhanced_score_notes(zscore, timings_divider=32)
        
        escore = TMIDIX.fix_escore_notes_durations(escore)
        
        cscore = TMIDIX.chordify_score([1000, escore])
        
        score = []
        
        pc = cscore[0]
        
        chords = []
        
        for c in cscore:
        
            pitches = sorted(set([e[4] for e in c]))
        
            if len(pitches) > 1:
                tones_chord = sorted(set([p % 12 for p in pitches]))
        
                if tones_chord not in TMIDIX.ALL_CHORDS_SORTED:
                    tones_chord = TMIDIX.check_and_fix_tones_chord(tones_chord, use_full_chords=False)
        
                chord_tok = TMIDIX.ALL_CHORDS_SORTED.index(tones_chord)+12
        
            if len(pitches) == 1:
                chord_tok = pitches[0] % 12
        
            score.append(chord_tok+128)
            chords.append(chord_tok+128)
        
            score.append(max(0, min(127, c[0][1]-pc[0][1])))
        
            for n in c:
                score.extend([max(1, min(127, n[4]))+461, max(1, min(127, n[2]))+589])
        
            pc = c

        return score, chords

    else:
        return [128]

def save_midi(tokens):
    
    """
    Convert token sequence back to a MIDI score and write it using TMIDIX.
    """

    song_f = []

    time = 0
    dur = 1
    vel = 90
    pitch = 60
    channel = 0
    patch = 0

    patches = [0] * 16

    for m in tokens:

        if 0 <= m < 128:
            time += m * 32

        elif 461 < m < 589:
            pitch = (m-461)

        elif 589 < m < 717:
            dur = (m-589) * 32
            song_f.append(['note', time, dur, 0, pitch, max(40, pitch), 0])

    if song_f is not None and song_f:

        song_f = TMIDIX.remove_duplicate_pitches_from_escore_notes(song_f)
    
        song_f = TMIDIX.fix_escore_notes_durations(song_f, 
                                                   min_notes_gap=0
                                                  )

        output_score = TMIDIX.humanize_velocities_in_escore_notes(song_f)
    
        now = datetime.datetime.now(PDT)
        ms4 = now.strftime("%f")[:4]  # first four digits of microseconds
        
        fname = (
            "Masked-Chordified-Piano-Transformer-Composition-"
            + now.strftime(f"%Y-%m-%d-%H-%M-%S-{ms4}")
        )

        os.makedirs(OUTPUT_MIDIS_DIR, exist_ok=True)

        output_fname = os.path.join(OUTPUT_MIDIS_DIR, fname)
        
        TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
            output_score,
            output_signature='Masked Chordified Piano Transformer',
            output_file_name=output_fname,
            track_name='Project Los Angeles',
            list_of_MIDI_patches=patches,
            verbose=False
        )
        return output_fname, output_score

    else:
        return None, None
        
# -----------------------------
# MUSIC GENERATION FUNCTIONS
# -----------------------------
@GPU
def generate_music(prime,
                   masked_gen,
                   chords,
                   num_gen_tokens,
                   num_gen_batches,
                   model_temperature,
                   model_top_p,
                   model_selector
                  ):
    
    """Generate music tokens given prime tokens and parameters."""

    if len(prime) >= 3072:
        prime = prime[-3072:]
    
    inputs = prime

    print(f'Will use {model_selector[0]}...')

    model = models_dict[model_selector[0]]

    model.to(MODEL_DEVICE)
    
    print("Generating...")
    inp = torch.LongTensor([inputs] * num_gen_batches).to(MODEL_DEVICE)

    if masked_gen:
        with ctx:
            out = model.generate_masked(
                inp,
                num_gen_tokens,
                filter_logits_fn=top_p,
                filter_kwargs={'thres': model_top_p},
                masked_token_ids=[i for i in range(128, 461) if i not in chords],
                temperature=model_temperature,
                return_prime=False,
                verbose=False
            )

    else:
        with ctx:
            out = model.generate(
                inp,
                num_gen_tokens,
                filter_logits_fn=top_p,
                filter_kwargs={'thres': model_top_p},
                temperature=model_temperature,
                return_prime=False,
                verbose=False
            )

    model.cpu()
            
    print("Done!")
    print_sep()
    return out.tolist()

def generate_music_and_state(input_midi,
                             masked_gen, 
                             num_prime_tokens, 
                             num_gen_tokens, 
                             model_temperature, 
                             model_top_p, 
                             final_composition, 
                             generated_batches, 
                             block_lines,
                             model_selector
                            ):
    
    """
    Generate tokens using the model, update the composition state, and prepare outputs.
    This function combines seed loading, token generation, and UI output packaging.
    """
    
    print_sep()
    print("Request start time:", datetime.datetime.now(PDT).strftime("%Y-%m-%d %H:%M:%S"))
    start_time = reqtime.time()

    print_sep()
    print('Requested model:', model_selector[0])
    
    if input_midi is not None:
        fn = os.path.basename(input_midi.name)
        fn1 = fn.split('.')[0]
        print('Input file name:', fn)

    print('Masked generation:', masked_gen)
    print('Num prime tokens:', num_prime_tokens)
    print('Num gen tokens:', num_gen_tokens)

    print('Model temp:', model_temperature)
    print('Model top p:', model_top_p)

    print_sep()

    chords = list(range(128, 461))
    
    # Load seed from MIDI if there is no existing composition.
    if not final_composition and input_midi is not None:
        final_composition, chords = load_midi(input_midi)

        if num_prime_tokens < 3072:
            final_composition = final_composition[:num_prime_tokens]
        
        midi_fname, midi_score = save_midi(final_composition)
        # Use the last note's time as a marker.
        last_nd_note = [e for e in midi_score if e[3] != 9]
        block_lines.append((last_nd_note[-1][1]+last_nd_note[-1][2]) // 1000 if final_composition else 0)

    if not final_composition and input_midi is None:
        final_composition = [128, 0]
            
    print_sep()
    print('Composition has', len(final_composition), 'tokens')
    print_sep()
    
    batched_gen_tokens = generate_music(final_composition,
                                        masked_gen,
                                        chords,
                                        num_gen_tokens,
                                        NUM_OUT_BATCHES,
                                        model_temperature,
                                        model_top_p,
                                        model_selector
                                       )
    
    output_batches = []
    for i, tokens in enumerate(batched_gen_tokens):
        preview_composition = final_composition
        preview_tokens = preview_composition[-PREVIEW_LENGTH:]

        plot_kwargs = {'plot_title': f'Batch # {i}', 'return_plt': True}
        
        if len(preview_composition) > PREVIEW_LENGTH:
            preview_score = save_midi(preview_tokens[:PREVIEW_LENGTH])[1]
            plot_kwargs['block_lines_times_list'] = [(preview_score[-1][1]+preview_score[-1][2]) // 1000]

        midi_fname, midi_score = save_midi(preview_tokens + tokens)
        midi_plot = TMIDIX.plot_ms_SONG(midi_score,
                                        **plot_kwargs
                                       )

        gradio_audio = midi_to_colab_audio(midi_fname + '.mid',
                                         soundfont_path=SOUNDFONT_PATH,
                                         sample_rate=AUDIO_SAMPLE_RATE,
                                         output_for_gradio=True)
        
        output_batches.append([(AUDIO_SAMPLE_RATE, gradio_audio), midi_plot, tokens, midi_fname + '.mid'])

    # Update generated_batches (for use by add/remove functions)
    generated_batches = batched_gen_tokens
    
    # Flatten outputs: states then audio and plots for each batch.
    outputs_flat = []
    for batch in output_batches:
        outputs_flat.extend([batch[0], batch[1], batch[3]])

    print("Request end time:", datetime.datetime.now(PDT).strftime("%Y-%m-%d %H:%M:%S"))
    print_sep()
    
    end_time = reqtime.time()
    execution_time = end_time - start_time
    
    print(f"Request execution time: {execution_time} seconds")
    print_sep()
        
    return [final_composition, generated_batches, block_lines] + outputs_flat

# -----------------------------
# BATCH HANDLING FUNCTIONS
# -----------------------------
def add_batch(batch_number, final_composition, generated_batches, block_lines):
    """Add tokens from the specified batch to the final composition and update outputs."""
    if generated_batches:
        final_composition.extend(generated_batches[batch_number])
        midi_fname, midi_score = save_midi(final_composition)
        last_nd_note = [e for e in midi_score if e[3] != 9]
        block_lines.append((last_nd_note[-1][1]+last_nd_note[-1][2]) // 1000 if final_composition else 0)
        midi_plot = TMIDIX.plot_ms_SONG(
            midi_score,
            plot_title='Masked Chordified Piano Transformer Composition',
            block_lines_times_list=block_lines[:-1],
            return_plt=True
        )
        gradio_audio = midi_to_colab_audio(midi_fname + '.mid',
                                         soundfont_path=SOUNDFONT_PATH,
                                         sample_rate=AUDIO_SAMPLE_RATE,
                                         output_for_gradio=True)
        print("Added batch #", batch_number)
        print_sep()
        return (AUDIO_SAMPLE_RATE, gradio_audio), midi_plot, midi_fname + '.mid', final_composition, generated_batches, block_lines
        
    else:
        return None, None, None, [], [], []

def remove_batch(batch_number, num_tokens, final_composition, generated_batches, block_lines):
    """Remove tokens from the final composition and update outputs."""
    if final_composition and len(final_composition) > num_tokens:
        final_composition = final_composition[:-num_tokens]
        if block_lines:
            block_lines.pop()
        midi_fname, midi_score = save_midi(final_composition)

        if midi_fname and midi_score:
            midi_plot = TMIDIX.plot_ms_SONG(
                midi_score,
                plot_title='Masked Chordified Piano Transformer Composition',
                block_lines_times_list=block_lines[:-1],
                return_plt=True
            )
            gradio_audio = midi_to_colab_audio(midi_fname + '.mid',
                                             soundfont_path=SOUNDFONT_PATH,
                                             sample_rate=AUDIO_SAMPLE_RATE,
                                             output_for_gradio=True)
            print("Removed batch #", batch_number)
            print_sep()
            return (AUDIO_SAMPLE_RATE, gradio_audio), midi_plot, midi_fname + '.mid', final_composition, generated_batches, block_lines

    return None, None, None, [], [], []

# -----------------------------
# MISC FUNCTIONS
# -----------------------------

def clear():
    """Clear outputs and reset state."""
    print_sep()
    print('Clear batch...')
    print_sep()
    return None, None, None, [], []

def reset(final_composition=[], generated_batches=[], block_lines=[]):
    """Reset composition state."""
    print_sep()
    print('Reset composition...')
    print_sep()
    return [], [], []

def update_state_from_dropdown(choice, state):
    """Store the dropdown value inside the global state list"""
    print_sep()
    print('Changed model from', state[0], 'to', choice)
    print_sep()
    state[0] = choice
    return state

Patch2number = TMIDIX.reverse_dict(TMIDIX.Number2patch)
Patch2number['Drums'] = 128

# -----------------------------
# GRADIO INTERFACE SETUP
# -----------------------------
with gr.Blocks() as demo:

    gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>Masked Chordified Piano Transformer</h1>")
    gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>Generate music with masked chordified Piano transformer</h1>")
    
    with gr.Row(elem_classes="duplicate-row"):
        gr.DuplicateButton(
            value="🤗 Duplicate 🤗",
            variant="huggingface",
            size="md",
            link="https://huggingface.co/spaces/projectlosangeles/Masked-Chordified-Piano-Transformer?duplicate=true",
            link_target="_blank"
        )
        
        gr.Button(
            value="❤️ Models ❤️",
            variant="huggingface",
            size="md",
            link="https://huggingface.co/asigalov61/Chordified-Piano-Transformer",
            link_target="_blank"
        )
        
        gr.Button(
            value="🦖 Dataset 🦖",
            variant="huggingface",
            size="md",
            link="https://huggingface.co/datasets/projectlosangeles/Godzilla-MIDI-Dataset",
            link_target="_blank"
        )



    # Global state variables for composition
    final_composition = gr.State([])
    generated_batches = gr.State([])
    block_lines = gr.State([])
    model_selector = gr.State([list(models_dict.keys())[0]])

    gr.Markdown("## Upload seed MIDI or simply click 'Generate' button for random output")
    
    input_midi = gr.File(label="Input MIDI", file_types=[".midi", ".mid", ".kar"])
    input_midi.upload(reset, [final_composition, generated_batches, block_lines],
                      [final_composition, generated_batches, block_lines])

    gr.Markdown("## Generation options")
    masked_gen = gr.Checkbox(value="True", label="Enable masked chords generation")
    
    num_prime_tokens = gr.Slider(16, 3072, value=512, step=1, label="Number of prime tokens")
    num_gen_tokens = gr.Slider(16, 1024, value=512, step=1, label="Number of tokens to generate")
    requested_model = gr.Dropdown(label="Model to use",
                                  choices=list(models_dict.keys()),
                                  value=list(models_dict.keys())[0]
                                 )
    model_temperature = gr.Slider(0.1, 1, value=0.9, step=0.01, label="Model temperature",
                                  info="Increase for more creative output, decrease for more repetitive output"
                                 )
    model_top_p = gr.Slider(0.1, 0.99, value=0.96, step=0.01, label="Model sampling top p value")
    generate_btn = gr.Button("Generate", variant="primary")

    gr.Markdown("## Batch Previews")
    outputs = [final_composition, generated_batches, block_lines]
    # Two outputs (audio and plot) for each batch
    for i in range(NUM_OUT_BATCHES):
        with gr.Tab(f"Batch # {i}"):
            audio_output = gr.Audio(label=f"Batch # {i} MIDI Audio", format=AUDIO_FORMAT)
            plot_output = gr.Plot(label=f"Batch # {i} MIDI Plot")
            midi_file = gr.File(label=f"Batch # {i} MIDI File")
            outputs.extend([audio_output, plot_output, midi_file])

    requested_model.change(
        fn=update_state_from_dropdown,
        inputs=[requested_model, model_selector],
        outputs=model_selector
    )
            
    generate_btn.click(
        generate_music_and_state,
        [input_midi,
         masked_gen, 
         num_prime_tokens, 
         num_gen_tokens, 
         model_temperature, 
         model_top_p,
         final_composition, 
         generated_batches, 
         block_lines,
         model_selector
        ],
        outputs
    )

    gr.Markdown("## Add/Remove Batch")
    batch_number = gr.Slider(0, NUM_OUT_BATCHES - 1, value=0, step=1, label="Batch number to add/remove")
    add_btn = gr.Button("Add batch", variant="primary")
    remove_btn = gr.Button("Remove batch", variant="stop")
    clear_btn = gr.ClearButton()

    final_audio_output = gr.Audio(label="Final MIDI audio", format=AUDIO_FORMAT)
    final_plot_output = gr.Plot(label="Final MIDI plot")
    final_file_output = gr.File(label="Final MIDI file")

    add_btn.click(
        add_batch,
        [batch_number, final_composition, generated_batches, block_lines],
        [final_audio_output, final_plot_output, final_file_output, final_composition, generated_batches, block_lines]
    )
    remove_btn.click(
        remove_batch,
        [batch_number, num_gen_tokens, final_composition, generated_batches, block_lines],
        [final_audio_output, final_plot_output, final_file_output, final_composition, generated_batches, block_lines]
    )
    clear_btn.click(clear, inputs=None,
                    outputs=[final_audio_output, final_plot_output, final_file_output, final_composition, block_lines])

# -----------------------------
# APP LAUNCHER
# -----------------------------
if __name__ == "__main__":
    demo.launch(
        mcp_server=RUNNING_IN_SPACE,   # MCP only on HF
        share=not RUNNING_IN_SPACE,    # Share only locally
        server_name="0.0.0.0",
        server_port=7860
    )