src/generation_service.py

"""MusicXML/MIDI ingestion, AMT accompaniment, and score export."""

from __future__ import annotations

import os
import random
import shutil
import tempfile
import zipfile
from io import BytesIO
from pathlib import Path
from typing import List, Literal, Sequence, Tuple

import mido
import numpy as np
import torch
from music21 import converter
from music21 import environment as m21_environment

from anticipation import ops
from anticipation.convert import events_to_midi, midi_to_events
from anticipation.sample import generate, generate_ar
from anticipation.tokenize import extract_instruments

Mode = Literal["anticipatory", "autoregressive"]


def set_generation_rng(seed: int) -> None:
    """Make torch.multinomial sampling reproducible for the same seed (local vs cloud, same stack)."""
    random.seed(seed)
    np.random.seed(seed % (2**32))
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        torch.mps.manual_seed(seed)
    try:
        import torch.backends.cudnn as cudnn

        cudnn.deterministic = True
        cudnn.benchmark = False
    except Exception:
        pass


_LILYPOND_CANDIDATES = (
    "/opt/homebrew/bin/lilypond",
    "/usr/local/bin/lilypond",
    "/usr/bin/lilypond",
)


def _find_lilypond() -> str | None:
    w = shutil.which("lilypond")
    if w:
        return w
    for p in _LILYPOND_CANDIDATES:
        if os.path.isfile(p) and os.access(p, os.X_OK):
            return p
    return None


def _configure_lilypond() -> str:
    """Point music21 at a LilyPond binary (required for lily.pdf)."""
    exe = _find_lilypond()
    if not exe:
        raise EnvironmentError(
            "LilyPond is not installed or not on your PATH, so PDF export cannot run. "
            "Install LilyPond: macOS: `brew install lilypond` · Ubuntu/Debian: `sudo apt install lilypond`. "
            "The Docker image for Hugging Face already includes LilyPond."
        )
    try:
        us = m21_environment.UserSettings()
        us["lilypondPath"] = exe
    except Exception:
        # UserSettings can fail in restricted environments; `which lilypond` may still suffice.
        pass
    return exe


class DummyTqdm:
    """Silence tqdm in anticipation.sample when running inside Streamlit."""

    def __init__(self, iterable=None, **_kwargs):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *args):
        return None

    def update(self, _n=1):
        pass


def _patch_tqdm():
    import anticipation.sample as sample_mod

    sample_mod.tqdm = lambda *args, **kwargs: DummyTqdm()


def _restore_tqdm():
    import anticipation.sample as sample_mod
    from tqdm import tqdm as real_tqdm

    sample_mod.tqdm = real_tqdm


def sniff_score_format(data: bytes) -> Literal["midi", "mxl", "musicxml"]:
    """
    Detect file type from bytes (do not rely on the browser filename).
    """
    if not data:
        raise ValueError("The file is empty. Choose a real score file, not a folder.")
    if len(data) < 4:
        raise ValueError("The file is too small to be a music score.")

    if data[:4] == b"MThd":
        return "midi"
    if data[:2] == b"PK":
        return "mxl"

    head = data[:4000].lstrip(b"\xef\xbb\xbf")
    if head.startswith(b"<?xml") or b"<score-partwise" in head or b"<score-timewise" in head:
        return "musicxml"

    raise ValueError(
        "This does not look like MusicXML or MIDI. "
        "Export from your notation app as **MusicXML** or **Standard MIDI** (not a project folder)."
    )


def _parse_musicxml_bytes_to_stream(data: bytes):
    """
    music21.converter.parse(BytesIO) sets an empty path string internally; cleanpath('')
    can resolve to the process cwd (a directory), so parseFile() opens the folder → EISDIR.
    Always parse MusicXML from a real temp file.
    """
    fd, path = tempfile.mkstemp(suffix=".musicxml")
    os.close(fd)
    try:
        with open(path, "wb") as f:
            f.write(data)
        return converter.parse(path, format="musicxml")
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


def _m21_stream_to_events_via_temp_midi(score_stream) -> List:
    """
    music21's MIDI writer uses open(path); BytesIO breaks or mis-resolves paths.
    Write to a real temp file, then read with mido.
    """
    fd, path = tempfile.mkstemp(suffix=".mid")
    os.close(fd)
    try:
        score_stream.write("midi", fp=path)
        mf = mido.MidiFile(path)
        return midi_to_events(mf)
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


def _music21_stream_from_midi_bytes(midi_bytes: bytes):
    fd, path = tempfile.mkstemp(suffix=".mid")
    try:
        with os.fdopen(fd, "wb") as tmp:
            tmp.write(midi_bytes)
        return converter.parse(path, format="midi")
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


def parse_upload_to_events(data: bytes, _filename: str | None = None) -> List:
    """Load MusicXML, MXL, or MIDI bytes into anticipation event tokens.

    Format is detected from the bytes (magic / XML header), not the filename.
    """
    kind = sniff_score_format(data)

    if kind == "midi":
        mf = mido.MidiFile(file=BytesIO(data))
        return midi_to_events(mf)

    if kind == "mxl":
        with zipfile.ZipFile(BytesIO(data)) as zf:
            xml_names = [n for n in zf.namelist() if n.endswith(".xml")]
            if not xml_names:
                raise ValueError("Compressed MusicXML (.mxl) has no XML inside.")
            xml_data = zf.read(xml_names[0])
        score_stream = _parse_musicxml_bytes_to_stream(xml_data)
        return _m21_stream_to_events_via_temp_midi(score_stream)

    # musicxml (text XML)
    if kind == "musicxml":
        score_stream = _parse_musicxml_bytes_to_stream(data)
        return _m21_stream_to_events_via_temp_midi(score_stream)

    raise ValueError("Unsupported format.")


def prepare_events_window(events: Sequence, clip_length_sec: float) -> List:
    """Clip to [0, clip_length_sec] and shift so the earliest event starts at t=0."""
    clipped = ops.clip(events, 0, clip_length_sec, clip_duration=False)
    if not clipped:
        return []
    t0 = ops.min_time(clipped, seconds=False)
    if t0 > 0:
        clipped = ops.translate(clipped, -t0, seconds=False)
    return clipped


def instrument_choices(events: Sequence) -> List[Tuple[int, int]]:
    """[(program, note_count), ...] descending by count."""
    inst = ops.get_instruments(events)
    return sorted(inst.items(), key=lambda x: -x[1])


def default_melody_programs(choices: Sequence[Tuple[int, int]]) -> List[int]:
    """Pick the busiest non-drum program as default melody."""
    non_drum = [(p, c) for p, c in choices if p != 128]
    if non_drum:
        return [non_drum[0][0]]
    if choices:
        return [choices[0][0]]
    return []


def validate_melody_vs_events(events: Sequence, melody_programs: Sequence[int]) -> None:
    """Ensure at least one locked melody note exists.

    Melody-only scores (one instrument, all notes locked) are valid: the model adds
    harmony on *other* instruments that were not in the original file.
    """
    _ev, ctr = extract_instruments(events, list(melody_programs))
    if not ctr:
        raise ValueError(
            "None of your selected parts appear in this file. Pick an instrument that has notes in the excerpt."
        )


def run_accompaniment(
    model,
    events: Sequence,
    melody_programs: Sequence[int],
    prompt_sec: float,
    clip_sec: float,
    top_p: float,
    mode: Mode,
    seed: int = 42,
) -> List:
    set_generation_rng(seed)

    events, controls = extract_instruments(events, list(melody_programs))
    prompt = ops.clip(events, 0, prompt_sec, clip_duration=False)

    _patch_tqdm()
    try:
        with torch.inference_mode():
            if mode == "anticipatory":
                gen = generate(model, prompt_sec, clip_sec, prompt, controls, top_p=top_p)
                out = ops.clip(ops.combine(gen, controls), 0, clip_sec)
            else:
                gen = generate_ar(model, prompt_sec, clip_sec, prompt, controls, top_p=top_p)
                out = ops.clip(gen, 0, clip_sec)
    finally:
        _restore_tqdm()

    return out


def events_to_midi_bytes(events: Sequence) -> bytes:
    mid = events_to_midi(events)
    bio = BytesIO()
    mid.save(file=bio)
    return bio.getvalue()


def events_to_musicxml_bytes(events: Sequence) -> bytes:
    score_stream = _music21_stream_from_midi_bytes(events_to_midi_bytes(events))
    fd, path = tempfile.mkstemp(suffix=".musicxml")
    os.close(fd)
    try:
        score_stream.write("musicxml", fp=path)
        with open(path, "rb") as f:
            return f.read()
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


def events_to_pdf_bytes(events: Sequence) -> bytes:
    """
    music21 needs the format `lily.pdf` (LilyPond), not plain `pdf`.
    """
    _configure_lilypond()
    score_stream = _music21_stream_from_midi_bytes(events_to_midi_bytes(events))
    tmpdir = tempfile.mkdtemp(prefix="amt_lily_")
    try:
        out_pdf = os.path.join(tmpdir, "score.pdf")
        written = score_stream.write("lily.pdf", fp=out_pdf)
        path = Path(written) if written is not None else Path(out_pdf)
        if not path.is_file():
            path = Path(out_pdf)
        if not path.is_file():
            raise RuntimeError(
                "LilyPond ran but no PDF was found. Check that `lilypond` works in a terminal "
                "(e.g. `lilypond --version`)."
            )
        return path.read_bytes()
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)