Blame for plus/plus_core.py

"""PLUS core: нотная модель, чтение/запись MIDI, темп и общие утилиты.

Всё остальное в пакете опирается на этот файл, поэтому здесь нет зависимостей
кроме pretty_midi и numpy. Ни одна функция не требует GPU.
"""
from __future__ import annotations

import json
import os
from dataclasses import dataclass, field, asdict
from typing import Iterable, Iterator, Sequence

import numpy as np

__version__ = "1.0"

DRUM_CHANNEL = 9


@dataclass
class Note:
    """Нота в секундах. pitch 0..127, velocity 1..127."""

    start: float
    dur: float
    pitch: int
    velocity: int = 80
    channel: int = 0
    track: str = "main"

    @property
    def end(self) -> float:
        return self.start + self.dur

    def shifted(self, delta: float) -> "Note":
        return Note(self.start + delta, self.dur, self.pitch, self.velocity, self.channel, self.track)

    def as_tuple(self) -> tuple[float, float, int, int, int]:
        return (self.start, self.dur, self.pitch, self.channel, self.velocity)


@dataclass
class Piece:
    """Набор нот плюс карта темпа и метаданные."""

    notes: list[Note] = field(default_factory=list)
    tempo_changes: list[tuple[float, float]] = field(default_factory=lambda: [(0.0, 120.0)])
    time_signature: tuple[int, int] = (4, 4)
    source: str = ""
    meta: dict = field(default_factory=dict)

    def copy(self) -> "Piece":
        return Piece([Note(**asdict(n)) for n in self.notes],
                     list(self.tempo_changes), self.time_signature, self.source, dict(self.meta))

    # --- производные величины -------------------------------------------------
    @property
    def duration(self) -> float:
        return max((n.end for n in self.notes), default=0.0)

    @property
    def bpm(self) -> float:
        return self.tempo_changes[0][1] if self.tempo_changes else 120.0

    def pitches(self) -> np.ndarray:
        return np.array([n.pitch for n in self.notes], dtype=int)

    def starts(self) -> np.ndarray:
        return np.array([n.start for n in self.notes], dtype=float)

    def velocities(self) -> np.ndarray:
        return np.array([n.velocity for n in self.notes], dtype=int)

    def stats(self) -> dict:
        if not self.notes:
            return {"notes": 0, "duration": 0.0}
        dur = self.duration or 1e-9
        return {
            "notes": len(self.notes),
            "duration": round(dur, 3),
            "notes_per_second": round(len(self.notes) / dur, 2),
            "pitch_min": int(min(n.pitch for n in self.notes)),
            "pitch_max": int(max(n.pitch for n in self.notes)),
            "velocity_min": int(min(n.velocity for n in self.notes)),
            "velocity_max": int(max(n.velocity for n in self.notes)),
            "bpm": round(self.bpm, 2),
        }


# --- чтение и запись ---------------------------------------------------------
def _pretty_midi():
    try:
        import pretty_midi  # noqa: F401
    except Exception as exc:  # pragma: no cover
        raise SystemExit("нужен pretty_midi: pip install pretty_midi (%s)" % exc)
    return pretty_midi


def load_midi(path: str) -> Piece:
    """MIDI-файл как Piece. Дорожки сохраняются по именам инструментов."""
    pm = _pretty_midi()
    midi = pm.PrettyMIDI(path)
    notes: list[Note] = []
    for instrument in midi.instruments:
        name = instrument.name or ("drums" if instrument.is_drum else "track%d" % instrument.program)
        for n in instrument.notes:
            notes.append(Note(float(n.start), max(1e-3, float(n.end - n.start)),
                              int(n.pitch), int(n.velocity) or 80,
                              int(instrument.program), name))
    notes.sort(key=lambda n: (n.start, n.pitch))
    times, bpms = midi.get_tempo_changes()
    tempo = [(float(t), float(b)) for t, b in zip(times, bpms)] or [(0.0, 120.0)]
    ts = midi.time_signature_changes[0] if midi.time_signature_changes else None
    piece = Piece(notes, tempo, (ts.numerator, ts.denominator) if ts else (4, 4),
                  source=os.path.basename(path))
    return piece


def save_midi(piece: Piece, path: str, split_tracks: bool = True) -> str:
    """Piece в MIDI. При split_tracks каждая дорожка идёт своим инструментом."""
    pm = _pretty_midi()
    midi = pm.PrettyMIDI(initial_tempo=piece.bpm or 120.0)
    groups: dict[str, list[Note]] = {}
    for n in piece.notes:
        groups.setdefault(n.track if split_tracks else "main", []).append(n)
    for index, (name, group) in enumerate(sorted(groups.items())):
        is_drum = name.lower() in ("drums", "drum")
        instrument = pm.Instrument(program=0 if is_drum else group[0].channel % 128,
                                   is_drum=is_drum, name=name)
        for n in sorted(group, key=lambda x: x.start):
            instrument.notes.append(pm.Note(int(max(1, min(127, int(n.velocity)))),
                                            int(max(0, min(127, int(n.pitch)))),
                                            float(n.start),
                                            float(n.end)))
        midi.instruments.append(instrument)
    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    midi.write(path)
    return path


# --- темп и сетка ------------------------------------------------------------
def explained_fraction(starts: Sequence[float], bpm: float, subdivisions: int = 4,
                       tolerance: float | None = None) -> float:
    """Доля онсетов, попадающих на сетку 1/subdivisions доли в пределах допуска."""
    onsets = np.asarray(list(starts), dtype=float)
    if onsets.size == 0 or bpm <= 0:
        return 0.0
    step = 60.0 / bpm / subdivisions
    tol = tolerance if tolerance is not None else max(0.012, step * 0.22)
    residual = np.abs(onsets - np.round(onsets / step) * step)
    return float((residual <= tol).mean())


def tempo_from_onsets(starts: Sequence[float], lo: float = 50.0, hi: float = 220.0,
                      center: float | None = None, spread: float = 1.25) -> float:
    """Темп, на чью сетку ложится больше всего онсетов.

    С бинарными онсетами чистые восьмые неотличимы от четвертей вдвое быстрее,
    поэтому без center выбирается самый МЕДЛЕННЫЙ темп с объяснением >= 0.92
    (он же самый правдоподобный для нотного текста), а ambiguity_of() показывает
    кратные варианты. С center поиск сужается до center*[1/spread, spread].
    """
    onsets = np.asarray(list(starts), dtype=float)
    if onsets.size < 4:
        return float(center or 120.0)
    low, high = (center / spread, center * spread) if center else (lo, hi)
    candidates = np.arange(max(20.0, low), min(400.0, high) + 0.01, 0.5)
    scored = [(explained_fraction(onsets, bpm), bpm) for bpm in candidates]
    best_fraction = max(fraction for fraction, _bpm in scored)
    if best_fraction < 0.5:
        return round(float(center or 120.0), 2)
    solid = [bpm for fraction, bpm in scored if fraction >= max(0.92, best_fraction - 0.03)]
    chosen = min(solid) if center is None else min(solid, key=lambda b: abs(b - center))
    return round(float(chosen), 2)


def ambiguity_of(starts: Sequence[float], bpm: float) -> list[dict]:
    """Кратные варианты темпа с их объяснением: честная оценка неоднозначности."""
    out = []
    for factor in (0.5, 1.0, 2.0, 3.0):
        candidate = bpm * factor
        if 20 <= candidate <= 400:
            out.append({"bpm": round(candidate, 2),
                        "explained": round(explained_fraction(starts, candidate), 3)})
    return out


def grid_times(piece: Piece, bpm: float | None = None, division: int = 4) -> np.ndarray:
    """Сетка долей (division долей на четверть) на всю длительность."""
    bpm = bpm or piece.bpm
    beat = 60.0 / bpm
    step = beat / division
    total = max(piece.duration + beat, step)
    return np.arange(0.0, total, step)


def nearest_grid(times: Iterable[float], grid: np.ndarray) -> np.ndarray:
    times = np.asarray(list(times), dtype=float)
    if grid.size == 0:
        return times
    index = np.clip(np.searchsorted(grid, times), 1, grid.size - 1)
    left, right = grid[index - 1], grid[index]
    return np.where(np.abs(times - left) <= np.abs(right - times), left, right)


# --- служебное ---------------------------------------------------------------
def step(number: int, total: int, text: str) -> None:
    """Единый формат прогресса: [этап] сообщение."""
    print("[%d/%d] %s" % (number, total, text), flush=True)


def json_default(obj):
    """numpy-скаляры в JSON: np.bool_ / np.float32 не сериализуются как есть."""
    if hasattr(obj, "item"):
        try:
            return obj.item()
        except Exception:
            pass
    if isinstance(obj, (set, tuple)):
        return list(obj)
    return str(obj)


def dump_json(payload: dict, path: str) -> str:
    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2,
                  default=json_default)
    return path


def pitch_name(pitch: int) -> str:
    names = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
    return "%s%d" % (names[int(pitch) % 12], int(pitch) // 12 - 1)