import spaces  # imported first, as ZeroGPU requires

import os
import sys
import math
import shutil
import tempfile
import pickle
from pathlib import Path

import torch
import numpy as np
import yaml
import gradio as gr
from pyharp import LabelList, MidiLabel, ModelCard, build_endpoint
from pyharp.tags import Subcategory
from huggingface_hub import hf_hub_download

sys.path.insert(0, str(Path(__file__).parent))
from dataset import load_vocabs
from models.variants import build_model

os.environ["TOKENIZERS_PARALLELISM"] = "true"

REPO_ID = "Itsuki-music/BACHI_Chord_Recognition"
CHECKPOINT_NAMES = {
    "Classical": "classical_film_kdec",
    "Pop": "pop909_film_kdec",
}
CHECKPOINT_FILES = ["best_model.pt", "config.yaml", "vocab.pkl"]

loaded_models = {}

def get_model(model_type: str):
    if model_type not in loaded_models:
        folder = CHECKPOINT_NAMES[model_type]
        ckpt_dir = Path(__file__).parent / "ckpts" / folder
        ckpt_dir.mkdir(parents=True, exist_ok=True)

        for fname in CHECKPOINT_FILES:
            dest = ckpt_dir / fname
            if not dest.exists():
                print(f"Downloading {folder}/{fname}...", flush=True)
                hf_hub_download(
                    repo_id=REPO_ID,
                    filename=f"{folder}/{fname}",
                    repo_type="dataset",
                    local_dir=Path(__file__).parent / "ckpts",
                    revision="7fd45be21b2b917d0a3fc349b6f48e9486f6ad19",
                )
                print(f"Downloaded {fname}.", flush=True)

        with open(ckpt_dir / "config.yaml", "r") as f:
            config = yaml.safe_load(f)

        vocabs = load_vocabs(str(ckpt_dir / "vocab.pkl"))
        use_key = (
            bool(config.get("use_key", False))
            or bool(config["training"].get("use_key", False))
            or bool(config["model"].get("use_key", False))
        )
        experiment = config["experiment"]

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Loading {model_type} model...", flush=True)
        model = build_model(experiment, config["model"], vocabs, use_key=use_key).to(device)
        model.load_state_dict(
            torch.load(ckpt_dir / "best_model.pt", map_location="cpu", weights_only=True)
        )
        model.eval()
        print(f"{model_type} model loaded.", flush=True)

        loaded_models[model_type] = (model, config, vocabs, use_key, device)

    return loaded_models[model_type]


def extract_pianoroll(score_path: Path, resolution: int = 12):
    import miditoolkit
    from music21 import converter, note as m21_note, chord as m21_chord

    notes_data = []
    suffix = score_path.suffix.lower()

    if suffix in {".mid", ".midi"}:
        midi = miditoolkit.MidiFile(str(score_path))
        tpb = midi.ticks_per_beat or 480
        for inst in midi.instruments:
            if inst.is_drum:
                continue
            for n in inst.notes:
                notes_data.append((n.pitch, n.start / tpb, n.end / tpb))
    else:
        sc = converter.parse(str(score_path))
        parts = list(sc.parts) if sc.hasPartLikeStreams() else [sc]
        for part in parts:
            inst = part.getInstrument()
            if inst:
                classes = inst.classes if hasattr(inst, "classes") else []
                if "Percussion" in classes or "Unpitched" in classes:
                    continue
            for el in part.flat.notes:
                dur = float(el.quarterLength)
                start = float(el.offset)
                if isinstance(el, m21_note.Note):
                    notes_data.append((el.pitch.midi, start, start + dur))
                elif isinstance(el, m21_chord.Chord):
                    for p in el.pitches:
                        notes_data.append((p.midi, start, start + dur))

    if not notes_data:
        return None

    notes_data.sort(key=lambda x: x[1])
    last_end = max(nd[2] for nd in notes_data)
    total_frames = math.ceil(last_end * resolution)
    pianoroll = np.zeros((88, total_frames), dtype=np.int8)
    for midi_pitch, start_b, end_b in notes_data:
        row = midi_pitch - 21
        if not (0 <= row < 88):
            continue
        s_f = max(0, math.floor(start_b * resolution))
        e_f = min(total_frames, math.ceil(end_b * resolution))
        pianoroll[row, s_f:e_f] = 1

    return pianoroll


def score_as_midi(score_path: Path) -> str:
    """The score as a MIDI file of its own, for the chord labels to be drawn over."""
    handle, midi_path = tempfile.mkstemp(suffix=".mid")
    os.close(handle)
    if score_path.suffix.lower() in {".mid", ".midi"}:
        shutil.copyfile(score_path, midi_path)
    else:
        from music21 import converter
        converter.parse(str(score_path)).write("midi", fp=midi_path)
    return midi_path


# The chord qualities of the vocabulary, as a symbol reads them
QUALITY_SUFFIX = {
    "M": "", "m": "m", "M7": "maj7", "m7": "m7", "D7": "7", "+": "+", "+7": "+7",
    "o": "dim", "o7": "dim7", "/o7": "m7b5", "mM7": "mMaj7", "sus2": "sus2", "sus4": "sus4",
    "other": "?",
}


def chord_symbol(root: str, quality: str, bass: str) -> str:
    if root == "N" or quality == "N":
        return "N.C."
    symbol = root + QUALITY_SUFFIX.get(quality, quality)
    if bass not in (root, "N", "PAD"):
        symbol += "/" + bass
    return symbol


def chord_labels(segments, midi_path: str) -> LabelList:
    """One label per chord segment, placed in seconds by the tempo map of the MIDI."""
    import miditoolkit

    midi = miditoolkit.MidiFile(midi_path)
    ticks_per_beat = midi.ticks_per_beat or 480
    mapping = midi.get_tick_to_time_mapping()
    final_tempo = midi.tempo_changes[-1].tempo if midi.tempo_changes else 120.0

    def seconds(beat: float) -> float:
        tick = int(round(beat * ticks_per_beat))
        if tick < len(mapping):
            return float(mapping[tick])
        # Past the last event of the file, time runs on at the final tempo
        last = len(mapping) - 1
        return float(mapping[last]) + (tick - last) / ticks_per_beat * 60.0 / final_tempo

    labels = LabelList()
    for start_beat, end_beat, root, quality, bass in segments:
        if root == "PAD":
            continue
        start, end = seconds(start_beat), seconds(end_beat)
        labels.append(MidiLabel(
            t=start,
            duration=max(0.0, end - start),
            label=chord_symbol(root, quality, bass),
            description=(f"Beats {start_beat:g} to {end_beat:g}: root {root}, quality {quality}, "
                         f"bass {bass}"),
        ))
    return labels


def predict_piece(pianoroll, model, config, vocabs, use_key, device):
    beat_resolution = config["model"]["beat_resolution"]
    label_resolution = config["model"]["label_resolution"]
    segment_len = config["model"]["n_beats"] * beat_resolution
    pr_to_label_ratio = beat_resolution // label_resolution
    comps_eval = ["root", "quality", "bass"] + (["key"] if use_key else [])

    n_frames = pianoroll.shape[0]
    segments, masks = [], []
    for i in range(0, n_frames, segment_len):
        seg = pianoroll[i : i + segment_len]
        orig_len = seg.shape[0]
        if orig_len < segment_len:
            seg = torch.cat([seg, torch.zeros(segment_len - orig_len, seg.shape[1])], dim=0)
        mask = torch.ones(segment_len, dtype=torch.bool)
        if orig_len < segment_len:
            mask[orig_len:] = False
        segments.append(seg)
        masks.append(mask)

    piece_preds = {k: [] for k in comps_eval + ["boundary"]}
    for i in range(0, len(segments), 16):
        batch_segs = torch.stack(segments[i : i + 16]).to(device)
        batch_masks = torch.stack(masks[i : i + 16]).to(device)
        with torch.no_grad():
            out = model.forward_infer(batch_segs, src_key_padding_mask=~batch_masks)
        for k in comps_eval + ["boundary"]:
            if k in out:
                piece_preds[k].append(out[k].detach().cpu())

    n_target = math.ceil(n_frames / pr_to_label_ratio)
    piece_pred_ids = {}
    for k, parts in piece_preds.items():
        if not parts:
            continue
        cat = torch.cat([p.reshape(-1) for p in parts], dim=0)
        piece_pred_ids[k] = cat[:n_target] if k == "boundary" else cat[:n_target].long()

    inv_root = {v: k for k, v in vocabs["root"].items()}
    inv_qual = {v: k for k, v in vocabs["quality"].items()}
    inv_bass = {v: k for k, v in vocabs["bass"].items()}

    valid_len = len(piece_pred_ids.get("root", []))
    if valid_len == 0:
        return []

    r_seq = piece_pred_ids["root"][:valid_len].tolist()
    q_seq = piece_pred_ids["quality"][:valid_len].tolist()
    b_seq = piece_pred_ids["bass"][:valid_len].tolist()

    # Each label token covers this many beats, and a segment runs until the chord changes
    beats_per_token = 1.0 / max(1, config["model"]["label_resolution"])
    segments = []
    cur_r, cur_q, cur_b = r_seq[0], q_seq[0], b_seq[0]
    cur_start = 0

    def close(end_token):
        segments.append((cur_start * beats_per_token, end_token * beats_per_token,
                         inv_root.get(cur_r, str(cur_r)), inv_qual.get(cur_q, str(cur_q)),
                         inv_bass.get(cur_b, str(cur_b))))

    for t in range(1, valid_len):
        if r_seq[t] != cur_r or q_seq[t] != cur_q or b_seq[t] != cur_b:
            close(t)
            cur_r, cur_q, cur_b = r_seq[t], q_seq[t], b_seq[t]
            cur_start = t
    close(valid_len)

    return segments


model_card = ModelCard(
    name="BACHI",
    description=(
        "Automatic chord recognition from symbolic music scores (MIDI or "
        "MusicXML). Returns the score as MIDI, with a chord label over each "
        "segment, aligned to the beat. Repository: "
        "https://github.com/AndyWeasley2004/BACHI_Chord_Recognition Paper: "
        "https://arxiv.org/abs/2510.06528"
    ),
    author="Mingyang Yao, Ke Chen, Shlomo Dubnov, Taylor Berg-Kirkpatrick",
    tags=[Subcategory.MUSIC_ANALYSIS, 'domain:music'],
)


def process_fn(input_file: str, model_type: str) -> tuple[str, LabelList]:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if model_type not in CHECKPOINT_NAMES:
        raise gr.Error(f"Unknown model type: {model_type}")
    if not input_file:
        raise gr.Error("Provide the score to analyze.")
    config = get_model(model_type)[1]

    score_path = Path(input_file)
    try:
        pianoroll_np = extract_pianoroll(score_path,
                                         resolution=config["model"]["beat_resolution"])
        midi_path = score_as_midi(score_path)
    except Exception as exc:
        raise gr.Error(f"Could not read the score: {exc}") from exc
    if pianoroll_np is None:
        raise gr.Error("Could not extract notes from the input file.")

    print(f"Processing with {model_type} model...", flush=True)
    segments = recognize(pianoroll_np, model_type)
    print("Done.", flush=True)
    if not segments:
        raise gr.Error("The model found no chords in this score.")

    return midi_path, chord_labels(segments, midi_path)


@spaces.GPU
def recognize(pianoroll_np, model_type: str) -> list:
    model, config, vocabs, use_key, device = get_model(model_type)
    pianoroll = torch.from_numpy(pianoroll_np.T).float()
    return predict_piece(pianoroll, model, config, vocabs, use_key, device)


# 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__":
    for _model_type in CHECKPOINT_NAMES:
        get_model(_model_type)


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.File(type="filepath",
                label="Input Score (.mid, .midi, .musicxml, .mxl, .xml)",
                file_types=[".mid", ".midi", ".musicxml", ".mxl", ".xml"],
            ).set_info((
                "A score to analyze. Chords are read from the notes, so a melody "
                "alone gives less to work with than a full texture."
            )),
            gr.Dropdown(
                choices=["Classical", "Pop"],
                value="Classical",
                label="Model Type",
                info=(
                    "Which trained model to use. The default was trained on the widest "
                    "range of material."
                ),
            ),
        ]
        output_components = [
            gr.File(
                type="filepath",
                file_types=[".mid"],
                label="Score as MIDI",
            ).set_info("The score, as MIDI, with the chord labels over it."),
            gr.JSON(label="Chord Labels").set_info(
                "One label per chord, from its first beat to the next change."
            ),
        ]
        app = build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

    print("Launching Gradio...", flush=True)
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)