from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

from pathlib import Path
from tempfile import mkdtemp

import gradio as gr
from pyharp import LabelList, MidiLabel, ModelCard, build_endpoint
from pyharp.tags import Subcategory

from runtime import generate, load_model

model_card = ModelCard(
    name="MuseMorphose",
    description=(
        "Transform piano music with controls for rhythmic intensity and "
        "polyphonicity. The MIDI is read as one piano part in 4/4, on a grid of sixteenth notes, "
        "with drums left out. The code is MIT, but the checkpoint is CC-BY-4.0, which requires "
        "attribution. Repository: https://github.com/YatingMusic/MuseMorphose "
        "Paper: https://arxiv.org/abs/2105.04090"
    ),
    author="Shih-Lun Wu and Yi-Hsuan Yang",
    tags=[Subcategory.EDITING, 'style-transfer', 'piano'],
)


def process_fn(
    midi_path,
    start_bar,
    bars,
    rhythm_shift,
    poly_shift,
    temperature,
    top_p,
    seed,
):
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not midi_path:
        raise gr.Error("Provide the MIDI to transform.")
    if int(bars) > 4:
        raise gr.Error("Choose at most four bars per request.")
    if not 0.5 <= temperature <= 2 or not 0 < top_p <= 1:
        raise gr.Error("Temperature or top-p is out of range.")

    reference, transformed, generated_remi, details = transform(
        midi_path, start_bar, bars, rhythm_shift, poly_shift, temperature, top_p, seed
    )

    # HARP reads a JSON output as labels, so the details are one label at the start of the result
    labels = LabelList()
    labels.append(MidiLabel(
        t=0.0,
        label=f"{details['generated_notes']} notes in {int(bars)} bars from bar {int(start_bar)}",
        description=f"{details['generated_events']} REMI events generated",
    ))

    return reference, transformed, generated_remi, labels


@spaces.GPU
def transform(midi_path, start_bar, bars, rhythm_shift, poly_shift, temperature, top_p, seed):
    try:
        # mkdtemp, not a context manager: that would delete the directory on
        # exit, taking the three files with it before the caller can read
        # them, which is what the block-cache calls here were working around.
        # Those also reached for the output component list, which exists only
        # under __main__ and so is undefined in the worker that imports this.
        reference, transformed, generated_remi, details = generate(
            midi_path,
            start_bar,
            int(bars),
            rhythm_shift,
            poly_shift,
            temperature,
            top_p,
            seed,
            Path(mkdtemp(prefix="musemorphose-")),
        )

    except Exception as exc:
        raise gr.Error(f"MuseMorphose generation failed: {exc}") from exc

    return reference, transformed, generated_remi, details


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is loaded there, before the first request rather than during it.
if __name__ != "__main__":
    load_model()


if __name__ == "__main__":
    with gr.Blocks(title="MuseMorphose", delete_cache=(3600, 86400)) as demo:
        gr.Markdown("## MuseMorphose piano transformation")
        inputs = [
            gr.File(
                type="filepath",
                file_types=[".mid", ".midi"],
                label="Input MIDI",
            ).harp_required(True).set_info(
                "Piano music to transform. Every pitched track is merged into one piano part, "
                "read in 4/4 on a grid of sixteenth notes, and drums are left out."
            ),
        ]
        inputs.extend(
            [
                gr.Number(value=0, minimum=0, maximum=9999, precision=0, label="Start bar",
                          info="The first bar to transform, counting from 0."),
                gr.Dropdown(
                    choices=["1", "2", "3", "4"], value="1", label="Number of bars",
                    info="How many bars to transform, counting from the start bar.",
                ),
                gr.Slider(-3, 3, step=1, value=0, label="Rhythmic intensity shift",
                          info="Steps by which to raise or lower each bar's rhythmic intensity, "
                               "on the model's scale of 0 to 7. Higher values give busier "
                               "rhythms."),
                gr.Slider(-3, 3, step=1, value=0, label="Polyphonicity shift",
                          info="Steps by which to raise or lower each bar's polyphony, on the "
                               "model's scale of 0 to 7. Higher values sound more notes at once."),
                gr.Slider(0.5, 2.0, step=0.1, value=1.2, label="Temperature",
                          info="Sampling randomness. Higher values give more varied but less "
                               "stable playing."),
                gr.Slider(0.1, 1.0, step=0.05, value=0.9, label="Top-p",
                          info="Samples only from the likeliest tokens whose probabilities add "
                               "up to this. Lower values stay closer to the likeliest playing."),
                gr.Number(
                    value=0, minimum=0, maximum=2_147_483_647, precision=0, label="Seed",
                    info="Use the same seed to reproduce a result.",
                ),
            ]
        )
        outputs = [
            gr.File(type="filepath", label="Reference MIDI", file_types=[".mid", ".midi"]).set_info((
                "The chosen bars as uploaded, rendered as MIDI for comparison with the "
                "transformed ones."
            )),
            gr.File(type="filepath", label="Transformed MIDI", file_types=[".mid", ".midi"]).set_info((
                "The restyled performance."
            )),
            gr.File(type="filepath", label="Generated REMI events", file_types=[".remi"]).set_info((
                "The raw REMI token sequence behind the MIDI, for inspection."
            )),
            gr.JSON(label="Generation details").set_info(
                "A label summarizing what was generated."
            ),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=inputs,
            output_components=outputs,
            process_fn=process_fn,
        )


    demo.queue(default_concurrency_limit=1, max_size=8).launch(pwa=True, show_error=True)