theme_generation/io.py

"""ABC and MIDI export helpers."""

from __future__ import annotations

import sys
from pathlib import Path

from .common import ABC_DURATION, COMMON_DURATION_BEATS, PC_TO_SHARP_NAME, PITCH_CLASS, Symbol


MUSES_SRC_PATH = Path("/Users/francoispachet/IdeaProjects/muses/src")
MUSES_VENV_LIB = Path("/Users/francoispachet/IdeaProjects/muses/.venv/lib")
# Generated symbols carry pitch class but not register; keep realization readable on treble staff.
DEFAULT_MIN_PITCH = 60
DEFAULT_MAX_PITCH = 84
DEFAULT_CENTER_PITCH = 67


def ensure_muses_import_path() -> None:
    if MUSES_SRC_PATH.exists():
        sys.path.insert(0, str(MUSES_SRC_PATH))
    if MUSES_VENV_LIB.exists():
        for candidate in MUSES_VENV_LIB.glob("python*/site-packages"):
            sys.path.insert(0, str(candidate))


def symbol_to_midi_pitch(
    symbol: Symbol,
    key_pc: int,
    previous_pitch: int | None,
    *,
    min_pitch: int = DEFAULT_MIN_PITCH,
    max_pitch: int = DEFAULT_MAX_PITCH,
    center_pitch: int = DEFAULT_CENTER_PITCH,
) -> int:
    target_pc = (key_pc + symbol.rpc) % 12
    if previous_pitch is None:
        base = 60 + target_pc
        if base < 60:
            base += 12
        return base
    candidates = [
        target_pc + 12 * octave
        for octave in range(2, 9)
        if min_pitch <= target_pc + 12 * octave <= max_pitch
    ]
    if not candidates:
        candidates = [target_pc + 12 * octave for octave in range(2, 9)]
    return min(candidates, key=lambda pitch: (abs(pitch - previous_pitch), abs(pitch - center_pitch)))


def midi_to_abc_note(pitch: int) -> str:
    pc = pitch % 12
    octave = pitch // 12 - 1
    name = PC_TO_SHARP_NAME[pc]
    accidental = "^" if "#" in name else ""
    letter = name[0]
    if octave >= 5:
        note = letter.lower()
        suffix = "'" * (octave - 5)
    elif octave == 4:
        note = letter
        suffix = ""
    else:
        note = letter
        suffix = "," * (4 - octave)
    return accidental + note + suffix


def sequence_to_pitches(sequence: tuple[Symbol, ...], key_pc: int) -> list[int]:
    pitches = []
    previous = None
    for symbol in sequence:
        pitch = symbol_to_midi_pitch(symbol, key_pc, previous)
        pitches.append(pitch)
        previous = pitch
    return pitches


def sequence_to_abc(sequence: tuple[Symbol, ...], key_name: str, title: str, sample_id: int, composer: str) -> str:
    key_pc = PITCH_CLASS[key_name]
    pitches = sequence_to_pitches(sequence, key_pc)
    tokens = [
        midi_to_abc_note(pitch) + ABC_DURATION.get(symbol.duration, "4")
        for pitch, symbol in zip(pitches, sequence)
    ]
    bars = []
    current = []
    beat_sum = 0.0
    for token, symbol in zip(tokens, sequence):
        current.append(token)
        beat_sum += COMMON_DURATION_BEATS.get(symbol.duration, 1.0)
        if beat_sum >= 4.0:
            bars.append(" ".join(current))
            current = []
            beat_sum = 0.0
    if current:
        bars.append(" ".join(current))
    body = "| ".join(bars) + "|"
    return "\n".join(
        [
            f"X:{sample_id}",
            f"T:{title}",
            f"C:{composer}",
            "M:4/4",
            "L:1/16",
            "Q:1/4=120",
            f"K:{key_name}",
            body,
            "",
        ]
    )


def write_midi(sequence: tuple[Symbol, ...], key_name: str, path: Path) -> None:
    try:
        from mido import Message, MetaMessage, MidiFile, MidiTrack, bpm2tempo
    except ImportError:
        muses_site = Path("/Users/francoispachet/IdeaProjects/muses/.venv/lib")
        for candidate in muses_site.glob("python*/site-packages"):
            sys.path.insert(0, str(candidate))
        from mido import Message, MetaMessage, MidiFile, MidiTrack, bpm2tempo

    key_pc = PITCH_CLASS[key_name]
    pitches = sequence_to_pitches(sequence, key_pc)
    ticks_per_beat = 480
    mid = MidiFile(ticks_per_beat=ticks_per_beat)
    track = MidiTrack()
    mid.tracks.append(track)
    track.append(MetaMessage("time_signature", numerator=4, denominator=4, time=0))
    track.append(MetaMessage("key_signature", key=key_name, time=0))
    track.append(MetaMessage("set_tempo", tempo=bpm2tempo(120), time=0))
    for pitch, symbol in zip(pitches, sequence):
        duration_ticks = max(1, int(round(COMMON_DURATION_BEATS.get(symbol.duration, 1.0) * ticks_per_beat)))
        track.append(Message("note_on", note=pitch, velocity=72, time=0))
        track.append(Message("note_off", note=pitch, velocity=0, time=duration_ticks))
    track.append(MetaMessage("end_of_track", time=0))
    mid.save(path)


def sequence_to_muses_piece(sequence: tuple[Symbol, ...], key_name: str, title: str, composer: str):
    """Convert a generated theme sequence to a MusES Piece for notation export."""

    ensure_muses_import_path()
    try:
        from muses.base.temporals import Piece, TemporalCollection
    except ImportError as exc:  # pragma: no cover - depends on sibling project.
        raise RuntimeError(
            "MusicXML export requires the sibling muses package at "
            "/Users/francoispachet/IdeaProjects/muses."
        ) from exc

    key_pc = PITCH_CLASS[key_name]
    pitches = sequence_to_pitches(sequence, key_pc)
    melody = TemporalCollection(name="theme", instrument="piano", program_change=0)
    start = 0.0
    for pitch, symbol in zip(pitches, sequence):
        duration = COMMON_DURATION_BEATS.get(symbol.duration)
        if duration is None:
            raise ValueError(f"Unsupported generated duration {symbol.duration!r}")
        melody.insert_note(pitch, start, duration, velocity=72, midi_channel=0)
        start += duration

    return Piece(
        name=title,
        title=title,
        composer=composer,
        melodies=[melody],
        ticks_per_beat=480,
        time_signature="4/4",
        key_signature=key_name,
        tempo=500000,
    )


def write_sequence_musicxml(sequence: tuple[Symbol, ...], key_name: str, path: Path, *, title: str, composer: str) -> None:
    ensure_muses_import_path()
    try:
        from muses.io.musicxml import write_musicxml as write_muses_musicxml
    except ImportError as exc:  # pragma: no cover - depends on sibling project.
        raise RuntimeError(
            "MusicXML export requires muses.io.musicxml in the sibling muses package."
        ) from exc

    piece = sequence_to_muses_piece(sequence, key_name, title, composer)
    write_muses_musicxml(piece, path)


def write_samples(
    sequences: list[tuple[Symbol, ...]],
    *,
    output_dir: Path,
    key_name: str,
    engine_name: str,
    write_abc: bool,
    write_musicxml_files: bool = False,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    for index, sequence in enumerate(sequences, start=1):
        title = f"{engine_name} {index:02d}"
        write_midi(sequence, key_name, output_dir / f"generated_{index:02d}.mid")
        if write_abc:
            abc = sequence_to_abc(sequence, key_name, title, index, engine_name)
            (output_dir / f"generated_{index:02d}.abc").write_text(abc, encoding="utf-8")
        if write_musicxml_files:
            write_musicxml_file = output_dir / f"generated_{index:02d}.musicxml"
            write_sequence_musicxml(sequence, key_name, write_musicxml_file, title=title, composer=engine_name)