#====================================================================
# https://huggingface.co/spaces/asigalov61/Godzilla-Piano-Transformer
#====================================================================

"""
Godzilla Piano Transformer Gradio App - Single Model, Simplified Version
Fast 807M 4k solo Piano music transformer trained on 1.14M+ MIDIs (2.7M+ samples)
Using only one model: "without velocity - 3 epochs"
"""

import os

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

import time as reqtime
import datetime
from pytz import timezone

import torch
import matplotlib.pyplot as plt
import gradio as gr
import spaces

from huggingface_hub import hf_hub_download
import TMIDIX
from midi_to_colab_audio import midi_to_colab_audio
from x_transformer_2_3_1 import TransformerWrapper, AutoregressiveWrapper, Decoder

# -----------------------------
# CONFIGURATION & GLOBALS
# -----------------------------
SEP = '=' * 70
PDT = timezone('US/Pacific')

MODEL_CHECKPOINT = 'Godzilla_Piano_Transformer_No_Velocity_Trained_Model_21113_steps_0.3454_loss_0.895_acc.pth'
SOUDFONT_PATH = 'SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2'
NUM_OUT_BATCHES = 12
PREVIEW_LENGTH = 120  # in tokens

# -----------------------------
# PRINT START-UP INFO
# -----------------------------
def print_sep():
    print(SEP)

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

# -----------------------------
# ENVIRONMENT & PyTorch Settings
# -----------------------------
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)

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

# -----------------------------
# MODEL INITIALIZATION
# -----------------------------
print_sep()
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 = 4096
PAD_IDX = 384

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

print_sep()
print("Loading model checkpoint...")
checkpoint = hf_hub_download(
    repo_id='asigalov61/Godzilla-Piano-Transformer',
    filename=MODEL_CHECKPOINT
)
model.load_state_dict(torch.load(checkpoint, map_location='cuda', weights_only=True))
model = torch.compile(model, mode='max-autotune')
print_sep()
print("Done!")
print("Model will use", dtype, "precision...")
print_sep()

model.cuda()
model.eval()

# -----------------------------
# HELPER FUNCTIONS
# -----------------------------
def render_midi_output(final_composition):
    """Generate MIDI score, plot, and audio from final composition."""
    fname, midi_score = save_midi(final_composition)
    time_val = midi_score[-1][1] / 1000  # seconds marker from last note
    midi_plot = TMIDIX.plot_ms_SONG(
        midi_score,
        plot_title='Godzilla Piano Transformer Composition',
        block_lines_times_list=[],
        return_plt=True
    )
    midi_audio = midi_to_colab_audio(
        fname + '.mid',
        soundfont_path=SOUDFONT_PATH,
        sample_rate=16000,
        output_for_gradio=True
    )
    return (16000, midi_audio), midi_plot, fname + '.mid', time_val

# -----------------------------
# MIDI PROCESSING FUNCTIONS
# -----------------------------
def load_midi(input_midi):
    """Process the input MIDI file and create a token sequence using without velocity logic."""
    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
    )[0]
    sp_escore_notes = TMIDIX.solo_piano_escore_notes(escore_notes)
    zscore = TMIDIX.recalculate_score_timings(sp_escore_notes)
    zscore = TMIDIX.augment_enhanced_score_notes(zscore, timings_divider=32)
    fscore = TMIDIX.fix_escore_notes_durations(zscore)
    cscore = TMIDIX.chordify_score([1000, fscore])
    
    score = []
    prev_chord = cscore[0]
    for chord in cscore:
        # Time difference token.
        score.append(max(0, min(127, chord[0][1] - prev_chord[0][1])))
        for note in chord:
            score.extend([
                max(1, min(127, note[2])) + 128,
                max(1, min(127, note[4])) + 256
            ])
        prev_chord = chord
    return score

def save_midi(tokens, batch_number=None):
    """Convert token sequence back to a MIDI score and write it using TMIDIX (without velocity).
       The output MIDI file name incorporates a date-time stamp.
    """
    song_events = []
    time_marker = 0
    duration = 0
    pitch = 0
    patches = [0] * 16

    for token in tokens:
        if 0 <= token < 128:
            time_marker += token * 32
        elif 128 <= token < 256:
            duration = (token - 128) * 32
        elif 256 <= token < 384:
            pitch = token - 256
            song_events.append(['note', time_marker, duration, 0, pitch, max(40, pitch), 0])
        # No velocity tokens are used.
    
    # Generate a time stamp using the PDT timezone.
    timestamp = datetime.datetime.now(PDT).strftime("%Y%m%d_%H%M%S")
    
    '''if batch_number is None:
        fname = f"Godzilla-Piano-Transformer-Music-Composition_{timestamp}"
    else:
        fname = f"Godzilla-Piano-Transformer-Music-Composition_{timestamp}_Batch_{batch_number}"'''

    fname = f"Godzilla-Piano-Transformer-Music-Composition"
    
    TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
        song_events,
        output_signature='Godzilla Piano Transformer',
        output_file_name=fname,
        track_name='Project Los Angeles',
        list_of_MIDI_patches=patches,
        verbose=False
    )
    return fname, song_events

# -----------------------------
# MUSIC GENERATION FUNCTION (Combined)
# -----------------------------
@spaces.GPU
def generate_music(prime, num_gen_tokens, num_mem_tokens, num_gen_batches, model_temperature):
    """Generate music tokens given prime tokens and parameters."""
    inputs = prime[-num_mem_tokens:] if prime else [0]
    print("Generating...")
    inp = torch.LongTensor([inputs] * num_gen_batches).cuda()
    with ctx:
        out = model.generate(
            inp,
            num_gen_tokens,
            temperature=model_temperature,
            return_prime=False,
            verbose=False
        )
    print("Done!")
    print_sep()
    return out.tolist()

def generate_music_and_state(input_midi, num_prime_tokens, num_gen_tokens, num_mem_tokens,
                             model_temperature, final_composition, generated_batches, block_lines):
    """
    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"))

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

    print('Num prime tokens:', num_prime_tokens)
    print('Num gen tokens:', num_gen_tokens)
    print('Num mem tokens:', num_mem_tokens)

    print('Model temp:', model_temperature)
    print('=' * 70)
    
    # Load seed from MIDI if there is no existing composition.
    if not final_composition and input_midi is not None:
        final_composition = load_midi(input_midi)[:num_prime_tokens]
        midi_fname, midi_score = save_midi(final_composition)
        # Use the last note's time as a marker.
        TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
            midi_score,
            output_signature='Godzilla Piano Transformer',
            output_file_name=midi_fname,
            track_name='Project Los Angeles',
            list_of_MIDI_patches=[0]*16,
            verbose=False
        )
        block_lines.append(midi_score[-1][1] / 1000 if final_composition else 0)
    
    batched_gen_tokens = generate_music(final_composition, num_gen_tokens, num_mem_tokens,
                                        NUM_OUT_BATCHES, model_temperature)
    
    output_batches = []
    for i, tokens in enumerate(batched_gen_tokens):
        preview_tokens = final_composition[-PREVIEW_LENGTH:]
        midi_fname, midi_score = save_midi(preview_tokens + tokens, batch_number=i)
        plot_kwargs = {'plot_title': f'Batch # {i}', 'return_plt': True}
        if len(final_composition) > PREVIEW_LENGTH:
            plot_kwargs['preview_length_in_notes'] = len([t for t in preview_tokens if t > 256])
        TMIDIX.Tegridy_ms_SONG_to_MIDI_Converter(
            midi_score,
            output_signature='Godzilla Piano Transformer',
            output_file_name=midi_fname,
            track_name='Project Los Angeles',
            list_of_MIDI_patches=[0]*16,
            verbose=False
        )
        midi_plot = TMIDIX.plot_ms_SONG(midi_score, **plot_kwargs)
        midi_audio = midi_to_colab_audio(midi_fname + '.mid',
                                         soundfont_path=SOUDFONT_PATH,
                                         sample_rate=16000,
                                         output_for_gradio=True)
        output_batches.append([(16000, midi_audio), midi_plot, tokens])
    
    # Update generated_batches (for use by add/remove functions)
    generated_batches = batched_gen_tokens
    
    print("Request end time:", datetime.datetime.now(PDT).strftime("%Y-%m-%d %H:%M:%S"))
    print_sep()
    
    # Flatten outputs: states then audio and plots for each batch.
    outputs_flat = []
    for batch in output_batches:
        outputs_flat.extend([batch[0], batch[1]])
    return [final_composition, generated_batches, block_lines] + outputs_flat

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

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

    gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>MuseCraft Piano Transformer</h1>")
    gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>Solo Piano music transformer for MuseCraft project</h1>")
    gr.HTML("""
        Check out <a href="https://github.com/MIDIAI/MuseCraft">MuseCraft project</a> on GitHub
        <p>
            <a href="https://huggingface.co/spaces/projectlosangeles/MuseCraft-Piano-Transformer?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 generation!
    """)

    # Global state variables for composition
    final_composition = gr.State([])

    gr.Markdown("## Upload seed MIDI or click 'Generate' for a 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("## Generate")
    num_prime_tokens = gr.Slider(15, 3072, value=3072, step=1, label="Number of prime tokens")
    num_gen_tokens = gr.Slider(15, 1024, value=512, step=1, label="Number of tokens to generate")
    num_mem_tokens = gr.Slider(15, 4096, value=4096, step=1, label="Number of memory tokens")
    model_temperature = gr.Slider(0.1, 1, value=0.9, step=0.01, label="Model temperature")
    generate_btn = gr.Button("Generate", variant="primary")

    gr.Markdown("## MIDI Output")
            
    generate_btn.click(
        generate_music_and_state,
        [input_midi, num_prime_tokens, num_gen_tokens, num_mem_tokens, model_temperature],
        generated_MIDI_file
    )

    generated_MIDI_file = gr.File(label="Generated MIDI file")

demo.launch()