#!/usr/bin/env python3
"""
mp3_to_yms.py — turns an audio file (mp3, wav, ogg, flac…) or a MIDI file (.mid) into a .yms song for
the Composer app of OpenOS Yaru (note blocks in OpenComputers).

What it does
  1. decodes the audio (soundfile, or ffmpeg as a fallback) to mono 22.05 kHz;
  2. estimates the tempo and the beat phase from onset strength;
  3. computes a pitch salience per semitone (harmonic sum over an STFT), so
     overtones do not show up as extra notes;
  4. on every grid step picks the voices:
       melody  — strongest salient note in the upper range,
       bass    — strongest note below MIDI 55,
       harmony — further strong notes (optional, more tracks);
     a note is written when it starts (onset), long notes can be re-struck;
  5. folds every track into the note block range F#3..F#5 (25 semitones)
     with the octave shift that keeps most notes in range;
  6. writes a Lua table readable by OpenOS `serialization.unserialize`
     and, optionally, a WAV preview of what the note blocks will play.

Note blocks only "pluck" (no sustain, one note per tick per block), so the
result is an arrangement of the song, not a copy of the recording.
Clear melodies (piano, guitar, chiptune, a solo voice) work best; dense mixes
give a busier, less accurate result: lower --tracks, raise --threshold.

Requirements:  py -m pip install numpy scipy soundfile
               (if soundfile cannot read mp3 on your system, install ffmpeg)
Examples:      py mp3_to_yms.py song.mp3 song.yms --preview song_preview.wav
               py mp3_to_yms.py song.mp3 --voices band
               py mp3_to_yms.py song.mp3 --voices melody:flute,harmony:bell,harmony:chime,bass::-12,kick,hat
"""
from __future__ import annotations

import argparse
import math
import os
import shutil
import subprocess
import sys

import numpy as np
from scipy import signal

SR = 22050
HOP = 256                      # ~11.6 ms analysis hop
NFFT = 4096                    # ~5.4 Hz bins: enough to separate semitones from ~E2 up
MIDI_LO, MIDI_HI = 36, 96      # analysed range C2..C7
NB_LO = 54                     # F#3 = note block pitch 0
NB_RANGE = 25                  # 0..24
# MIDI note of block pitch 0 for each instrument (Minecraft: bass/didgeridoo two
# octaves down, guitar one down, flute/cow bell one up, bell/chime/xylophone two up)
INST_BASE = {"harp": 54, "pling": 54, "banjo": 54, "bit": 54, "iron_xylophone": 54, "basedrum": 54, "snare": 54, "hat": 54,
             "bass": 30, "didgeridoo": 30, "guitar": 42, "flute": 66, "cow_bell": 66, "bell": 78, "chime": 78, "xylophone": 78}
INSTRUMENTS = ["harp", "basedrum", "snare", "hat", "bass", "flute", "bell", "guitar", "chime",
               "xylophone", "iron_xylophone", "cow_bell", "didgeridoo", "bit", "banjo", "pling"]


# ---------------------------------------------------------------- audio
def load_audio(path: str, start: float, duration: float | None) -> np.ndarray:
    data = None
    try:
        import soundfile as sf
        y, sr = sf.read(path, always_2d=True, dtype="float32")
        data = (y.mean(axis=1), sr)
    except Exception as e:  # noqa: BLE001 - any decoder problem -> try ffmpeg
        if not shutil.which("ffmpeg"):
            sys.exit(f"cannot decode {path}: {e}\ninstall ffmpeg or a soundfile build with mp3 support")
    if data is None:
        cmd = ["ffmpeg", "-v", "error", "-i", path, "-ac", "1", "-ar", str(SR), "-f", "f32le", "-"]
        raw = subprocess.run(cmd, check=True, stdout=subprocess.PIPE).stdout
        y, sr = np.frombuffer(raw, dtype=np.float32), SR
    else:
        y, sr = data
    if sr != SR:
        g = math.gcd(int(sr), SR)
        y = signal.resample_poly(y, SR // g, int(sr) // g).astype(np.float32)
    a = int(start * SR)
    b = len(y) if duration is None else min(len(y), a + int(duration * SR))
    y = y[a:b]
    peak = np.max(np.abs(y)) if len(y) else 0
    if peak > 0:
        y = y / peak
    return y


def stft_mag(y: np.ndarray, nfft: int = NFFT) -> np.ndarray:
    # centred frames: frame k is centred at k * HOP samples, for every window size
    _, _, z = signal.stft(y, fs=SR, window="hann", nperseg=nfft, noverlap=nfft - HOP, boundary="zeros", padded=True)
    return np.abs(z).astype(np.float32)          # (freq bins, frames)


# ---------------------------------------------------------------- tempo
def onset_envelope(mag: np.ndarray) -> np.ndarray:
    """spectral flux of a short-window STFT (accurate timing)"""
    logm = np.log1p(100 * mag)
    flux = np.maximum(0, np.diff(logm, axis=1)).sum(axis=0)
    flux = np.concatenate([[0], flux])
    flux = flux - signal.medfilt(flux, 31)
    return np.maximum(flux, 0)


def estimate_tempo(env: np.ndarray, lo: float = 60, hi: float = 200) -> float:
    fps = SR / HOP
    e = env - env.mean()
    ac = np.correlate(e, e, mode="full")[len(e) - 1:]
    best, best_score = 120.0, -1e9
    for bpm in np.arange(lo, hi, 0.5):
        lag = fps * 60 / bpm
        score = 0.0
        for k, w in ((1, 1.0), (2, 0.5), (4, 0.25)):   # beat + its multiples
            i = int(round(lag * k))
            if i < len(ac):
                score += w * ac[i]
        score *= math.exp(-0.5 * (math.log2(bpm / 120) / 0.9) ** 2)  # prefer ~120 bpm
        if score > best_score:
            best, best_score = bpm, score
    return float(best)


def beat_phase(env: np.ndarray, step_frames: float) -> float:
    """grid offset (frames) that puts the most onset energy on grid points"""
    best, best_v = 0.0, -1.0
    n = int(step_frames)
    for off in range(max(1, n)):
        idx = np.arange(off, len(env), step_frames).astype(int)
        v = env[idx].sum()
        if v > best_v:
            best, best_v = float(off), v
    return best


# ---------------------------------------------------------------- source separation (optional)
def separate(path: str, stem: str) -> str:
    """isolate the melody with Demucs (py -m pip install demucs): 'vocals' = sung
    melody, 'other' = everything but vocals/drums/bass (lead instruments).
    Returns the path of the separated WAV (cached next to the input)."""
    import tempfile
    out = os.path.join(tempfile.gettempdir(), "mp3_to_yms_demucs")
    name = os.path.splitext(os.path.basename(path))[0]
    wav = os.path.join(out, "htdemucs", name, f"{stem}.wav")
    if os.path.exists(wav):
        return wav
    cmd = [sys.executable, "-m", "demucs", "-n", "htdemucs", "-o", out]
    if stem == "vocals":
        cmd += ["--two-stems", "vocals"]
    cmd.append(path)
    print("Separate : running Demucs (first run downloads the model, a few minutes on CPU)…")
    r = subprocess.run(cmd)
    if r.returncode != 0 or not os.path.exists(wav):
        sys.exit("Demucs failed; install it with: py -m pip install demucs   (or drop --separate)")
    return wav


# ---------------------------------------------------------------- neural pitch (optional)
def find_model(path: str | None) -> str | None:
    """Basic Pitch (Spotify, Apache-2.0) ONNX model: --model, next to this script, or the installed package"""
    cands = [path] if path else []
    here = os.path.dirname(os.path.abspath(__file__))
    cands += [os.path.join(here, "nmp.onnx")]
    try:
        import basic_pitch  # noqa: F401
        cands.append(os.path.join(os.path.dirname(basic_pitch.__file__), "saved_models", "icassp_2022", "nmp.onnx"))
    except Exception:  # noqa: BLE001
        pass
    for c in cands:
        if c and os.path.exists(c):
            return c
    return None


def nn_posteriors(y: np.ndarray, model: str, n_frames: int) -> tuple[np.ndarray, np.ndarray]:
    """runs Basic Pitch over 2-second windows (30 frames overlap, as the reference code);
    returns (note, onset) posteriors resampled to our semitone rows MIDI_LO..MIDI_HI
    and our frame grid (both use a 256-sample hop at 22.05 kHz)"""
    import onnxruntime as ort
    sess = ort.InferenceSession(model, providers=["CPUExecutionProvider"])
    n_samples, olap_frames = 2 * SR - 256, 30
    olap = olap_frames * 256
    hop = n_samples - olap
    audio = np.concatenate([np.zeros(olap // 2, np.float32), y.astype(np.float32)])
    wins = []
    for i in range(0, len(audio), hop):
        w = audio[i:i + n_samples]
        if len(w) < n_samples:
            w = np.pad(w, (0, n_samples - len(w)))
        wins.append(w)
    notes, onsets = [], []
    inp = sess.get_inputs()[0].name
    names = [o.name for o in sess.get_outputs()]
    note_name = next(n for n in names if n.endswith(":1"))
    onset_name = next(n for n in names if n.endswith(":2"))
    for k in range(0, len(wins), 16):
        batch = np.stack(wins[k:k + 16])[..., None]
        out = sess.run([note_name, onset_name], {inp: batch})
        notes.append(out[0])
        onsets.append(out[1])
    h = olap_frames // 2
    def unwrap(parts):
        a = np.concatenate(parts)[:, h:-h, :]
        a = a.reshape(-1, a.shape[-1])
        return a[:int(len(y) * SR / 256 / SR) + 1]
    note, onset = unwrap(notes), unwrap(onsets)
    def to_rows(a):
        rows = np.zeros((MIDI_HI - MIDI_LO + 1, n_frames), np.float32)
        m = min(n_frames, a.shape[0])
        for midi in range(MIDI_LO, MIDI_HI + 1):
            rows[midi - MIDI_LO, :m] = a[:m, midi - 21]
        return rows
    return to_rows(note), to_rows(onset)


# ---------------------------------------------------------------- pitch
def salience(mag: np.ndarray, nfft: int = NFFT) -> np.ndarray:
    """harmonic-sum pitch salience, one row per MIDI note in MIDI_LO..MIDI_HI"""
    freqs = np.fft.rfftfreq(nfft, 1 / SR)
    comp = np.sqrt(mag)                       # compress dynamics
    out = np.zeros((MIDI_HI - MIDI_LO + 1, mag.shape[1]), np.float32)
    for m in range(MIDI_LO, MIDI_HI + 1):
        f0 = 440.0 * 2 ** ((m - 69) / 12)
        acc = np.zeros(mag.shape[1], np.float32)
        for h, w in ((1, 1.0), (2, 0.6), (3, 0.45), (4, 0.3), (5, 0.2)):
            f = f0 * h
            if f > SR / 2 - 50:
                break
            # band +-35 cents around the partial
            lo, hi = f * 2 ** (-0.35 / 12), f * 2 ** (0.35 / 12)
            band = (freqs >= lo) & (freqs <= hi)
            if not band.any():
                band = np.zeros_like(band)
                band[np.argmin(np.abs(freqs - f))] = True
            acc += w * comp[band].max(axis=0)
        out[m - MIDI_LO] = acc
    # suppress sub-octave ghosts: a note whose octave above is much stronger is an artefact
    for i in range(out.shape[0] - 12):
        out[i] = np.where(out[i + 12] > 1.6 * out[i], out[i] * 0.4, out[i])
    return out


# ---------------------------------------------------------------- transcription
def fold(notes: list[int], base: int = NB_LO) -> int:
    """octave shift that puts most notes into the instrument's 25-semitone range"""
    best, best_score = 0, -1e9
    for sh in range(-48, 49, 12):
        inside = sum(1 for n in notes if base <= n + sh < base + NB_RANGE)
        # keep the original octave unless moving gains clearly more notes (>10%)
        score = inside - 0.1 * len(notes) * abs(sh) / 12
        if score > best_score:
            best, best_score = sh, score
    return best


def to_block(m: int, base: int = NB_LO) -> int:
    p = m - base
    while p < 0:
        p += 12
    while p >= NB_RANGE:
        p -= 12
    return p


PITCHED = ("melody", "alto", "tenor", "bass", "harmony")
DRUMS = {"kick": ((30, 120), 3), "snare": ((1500, 5000), 10), "hat": ((7000, 11000), 18)}
DEFAULT_INST = {"melody": "harp", "alto": "flute", "tenor": "guitar", "bass": "bass", "harmony": ["bell", "chime", "xylophone", "pling", "banjo"],
                "kick": "basedrum", "snare": "snare", "hat": "hat"}
PRESETS = {"1": "melody", "2": "melody,bass", "3": "melody,bass,harmony",
           "piano": "melody,harmony,harmony,bass", "satb": "melody,alto,tenor,bass",
           "band": "melody,bass,harmony,kick,snare,hat", "full": "melody,alto,tenor,bass,harmony,harmony,kick,snare,hat",
           "drums": "kick,snare,hat"}


def parse_voices(spec: str) -> list[dict]:
    """'melody:flute,bass::-12,harmony,kick' -> [{role, inst, shift, name}]"""
    spec = PRESETS.get(spec, spec)
    out, count = [], {}
    for item in spec.split(","):
        item = item.strip()
        if not item:
            continue
        parts = item.split(":")
        role = parts[0].lower()
        if role not in PITCHED and role not in DRUMS:
            sys.exit(f"unknown voice {role!r}; use: {', '.join(PITCHED + tuple(DRUMS))} or a preset: {', '.join(PRESETS)}")
        count[role] = count.get(role, 0) + 1
        inst = parts[1] if len(parts) > 1 and parts[1] else None
        if inst is None and role in DRUMS:
            inst = DEFAULT_INST[role]
        if inst is not None and inst not in INSTRUMENTS:
            sys.exit(f"unknown instrument {inst!r}; use one of: {', '.join(INSTRUMENTS)}")
        shift = int(parts[2]) if len(parts) > 2 and parts[2] else 0
        name = role if count[role] == 1 else f"{role} {count[role]}"
        out.append({"role": role, "inst": inst, "shift": shift, "name": name, "nth": count[role]})
    if not out:
        sys.exit("no voices given")
    return out


_PERC = {}


def percussive(y: np.ndarray) -> np.ndarray:
    """percussive part of a short-window spectrogram (median-filter HPSS):
    drums are smooth across frequency, notes are smooth across time"""
    key = id(y)
    if key not in _PERC:
        from scipy.ndimage import median_filter
        mag = stft_mag(y, 1024)
        harm = median_filter(mag, size=(1, 17))
        perc = median_filter(mag, size=(17, 1))
        _PERC[key] = mag * (perc > 1.5 * harm)
    return _PERC[key]


def band_flux(y: np.ndarray, bands) -> np.ndarray:
    """onset strength of a frequency band of the percussive part"""
    mag = percussive(y)
    freqs = np.fft.rfftfreq(1024, 1 / SR)
    sel = np.zeros(len(freqs), bool)
    for k in range(0, len(bands), 2):
        sel |= (freqs >= bands[k]) & (freqs < bands[k + 1])
    e = np.log1p(100 * mag[sel]).sum(axis=0)
    flux = np.maximum(0, np.diff(e, prepend=e[:1]))
    return flux / (np.percentile(flux, 98) + 1e-9)


def make_grid(sal: np.ndarray, env: np.ndarray, bpm: float, spb: int) -> dict:
    fps = SR / HOP
    step_frames = fps * 60 / bpm / spb
    phase = beat_phase(env, step_frames)
    n_steps = int((sal.shape[1] - phase) / step_frames)
    steps = np.zeros((sal.shape[0], n_steps), np.float32)
    for s in range(n_steps):
        a = int(phase + (s + 0.25) * step_frames)
        b = max(a + 1, int(phase + (s + 1.25) * step_frames))
        steps[:, s] = sal[:, a:min(b, sal.shape[1])].mean(axis=1) if a < sal.shape[1] else 0
    loud = steps.max(axis=0) / (np.percentile(steps.max(axis=0), 95) + 1e-9)
    return {"step_frames": step_frames, "phase": phase, "n": n_steps, "steps": steps, "loud": loud}


def peaks(col: np.ndarray, thr: float) -> list[int]:
    """salient semitones of one step, without octave / octave+fifth ghosts of lower notes"""
    cand, kept = [], set()
    for i in range(1, len(col) - 1):
        if col[i] >= thr * 0.6 and col[i] >= col[i - 1] and col[i] >= col[i + 1]:
            if i >= 12 and (i - 12) in kept and col[i - 12] >= 0.8 * col[i]:
                continue
            if i >= 19 and (i - 19) in kept and col[i - 19] >= col[i]:
                continue
            kept.add(i)
            if col[i] >= thr:
                cand.append(i)
    return cand


def track_line(steps: np.ndarray, loud: np.ndarray, lo: int, hi: int, args, jump: float = 0.02,
               silence: float | None = None, avoid: list | None = None) -> list[int | None]:
    """Viterbi tracking of one continuous line (melody or bass) inside MIDI lo..hi:
    each step either holds a note or rests; the score rewards salience measured
    *within this range* (so a loud bass does not hide the melody) and penalises
    jumps, so the line follows the tune instead of jumping between instruments."""
    lo = max(lo, MIDI_LO + 1)
    a, b = lo - MIDI_LO, hi - MIDI_LO + 1
    sub = steps[a:b]
    n = sub.shape[1]
    if n == 0:
        return []
    norm = sub / (np.percentile(sub.max(axis=0), 90) + 1e-9)          # range-local loudness
    # a note that just started is more likely the tune than one still ringing
    rise = np.maximum(0, norm - np.concatenate([norm[:, :1], norm[:, :-1]], axis=1))
    states = sub.shape[0] + 1                                         # last state = rest
    rest_score = args.threshold * 0.55
    dist = np.abs(np.arange(sub.shape[0])[:, None] - np.arange(sub.shape[0])[None, :])
    trans = -jump * dist - np.where(dist > 7, 0.08, 0) - np.where(dist > 12, 0.15, 0)
    score = np.full(states, -1e9)
    back = np.zeros((n, states), np.int32)
    for s in range(n):
        col = norm[:, s]
        # only real peaks of this range can be notes (ghost-free, relative to the range max)
        if avoid is not None and avoid[s] is not None:
            # overtones of the bass note (octave, octave+fifth, two octaves) are not melody
            col = col.copy()
            for h in (12, 19, 24, 28):
                k = avoid[s] + h - lo
                if 0 <= k < len(col):
                    col[k] *= 0.35
        rel = col / (col.max() + 1e-9)
        cand = set(peaks(rel, args.threshold * 0.8)) if loud[s] >= (silence if silence is not None else args.silence) else set()
        emit = np.full(states, -1e9)
        for i in cand:
            emit[i] = col[i] + 0.8 * rise[i, s]
        emit[-1] = rest_score
        if s == 0:
            new = emit.copy()
            back[s] = states - 1
        else:
            new = np.full(states, -1e9)
            prev_notes = score[:-1]
            # to a note: from a note (with jump cost) or from rest (free)
            via_note = prev_notes[:, None] + trans                      # [from, to]
            best_from = via_note.argmax(axis=0)
            best_val = via_note.max(axis=0)
            from_rest = score[-1]
            use_rest = from_rest >= best_val
            new[:-1] = np.where(use_rest, from_rest, best_val) + emit[:-1]
            back[s, :-1] = np.where(use_rest, states - 1, best_from)
            # to rest: from anything
            k = int(score.argmax())
            new[-1] = score[k] + emit[-1]
            back[s, -1] = k
        score = new
    path = [0] * n
    k = int(score.argmax())
    for s in range(n - 1, -1, -1):
        path[s] = k
        k = back[s, k]
    return [None if p == states - 1 else lo + p for p in path]


def attack_steps(y, env, bpm, spb, n_steps, sustained, onset_post=None):
    """per-step salience of what *started* in the step: harmonic sum over the
    positive spectral difference (now minus ~90 ms ago). Held chords, pads and
    the bass hardly change, a newly played melody note does, so this separates
    the tune from the accompaniment much better than the plain spectrum."""
    if onset_post is not None:
        att = onset_post
    else:
        mag = stft_mag(y, NFFT)
        k = 8
        diff = np.maximum(0, mag - np.concatenate([np.repeat(mag[:, :1], k, axis=1), mag[:, :-k]], axis=1))
        att = salience(diff, NFFT)
    fps = SR / HOP
    step_frames = fps * 60 / bpm / spb
    phase = beat_phase(env, step_frames)
    out = np.zeros((att.shape[0], n_steps), np.float32)
    for s in range(n_steps):
        a = max(0, int(phase + (s - 0.3) * step_frames))
        b = max(a + 1, int(phase + (s + 0.7) * step_frames))
        seg = att[:, a:min(b, att.shape[1])]
        if seg.shape[1]:
            out[:, s] = seg.max(axis=1)
    na = out / (np.percentile(out.max(axis=0), 90) + 1e-9)
    ns = sustained / (np.percentile(sustained.max(axis=0), 90) + 1e-9)
    return na + 0.3 * ns[:, :n_steps]


def transcribe(sal: np.ndarray, env: np.ndarray, bpm: float, spb: int, args, voices_spec, y) -> tuple[dict, int, np.ndarray]:
    g = make_grid(sal, env, bpm, spb)
    step_frames, phase, n_steps, steps, loud = g["step_frames"], g["phase"], g["n"], g["steps"], g["loud"]
    def per_step(sig):
        v = np.zeros(n_steps, np.float32)
        for s in range(n_steps):
            a = int(phase + s * step_frames)
            v[s] = sig[a:int(phase + (s + 1) * step_frames) + 1].max() if a < len(sig) else 0
        return v
    env_steps = per_step(env)
    env_steps = env_steps / (np.percentile(env_steps, 95) + 1e-9)

    events = {v["name"]: [] for v in voices_spec}
    last: dict[str, tuple] = {}
    prevcol = np.zeros(steps.shape[0], np.float32)
    thr = args.threshold
    pitched = [v for v in voices_spec if v["role"] in PITCHED]

    def emit(name, i, s, col, onset, bass):
        m = i + MIDI_LO
        rise = col[i] > prevcol[i] * 1.25 + 0.05 or (s > 0 and steps[i, s] > steps[i, s - 1] * 1.12)
        prev = last.get(name)
        if prev is None or prev[0] != m:
            if not (prev and prev[0] == m and s - prev[1] <= 1):
                events[name].append((s, m)); last[name] = (m, s)
        elif s - prev[1] >= 2 and ((onset or bass) and rise or (args.restrike and s - prev[1] >= args.restrike)):
            events[name].append((s, m)); last[name] = (m, s)

    roles = {v["role"] for v in pitched}
    # the bass line first (long window: low notes need frequency resolution);
    # its overtones are then ignored by the melody tracker, which uses a shorter
    # window (2048 samples = 93 ms) so fast notes do not smear into each other
    bass_line = track_line(steps, loud, 28, 59, args, jump=0.04)
    mel_steps = steps
    if "melody" in roles:
        my = getattr(args, "melody_audio", None)
        if my is not None:
            msal = getattr(args, "mel_note", None)
            if msal is None:
                msal = salience(stft_mag(my))
            ms = make_grid(msal, env, bpm, spb)["steps"]
            ms = ms[:, :n_steps] if ms.shape[1] >= n_steps else np.pad(ms, ((0, 0), (0, n_steps - ms.shape[1])))
            mel_steps = attack_steps(my, env, bpm, spb, n_steps, ms, getattr(args, "onset_post", None))
        else:
            mel_steps = attack_steps(y, env, bpm, spb, n_steps, steps, getattr(args, "onset_post", None))
    lines = {"bass": bass_line if "bass" in roles else [None] * n_steps,
             "melody": track_line(mel_steps, loud, args.mel_lo, 96, args, avoid=bass_line) if "melody" in roles else [None] * n_steps}
    for s in range(n_steps):
        if loud[s] < args.silence or not pitched:
            prevcol = steps[:, s] * 0
            continue
        col = steps[:, s] / (steps[:, s].max() + 1e-9)
        cand = peaks(col, thr)
        if not cand and lines["melody"][s] is None and lines["bass"][s] is None:
            prevcol = col
            continue
        by_pitch = sorted(cand, reverse=True)           # high -> low
        used = set()
        onset = env_steps[s] > args.onset
        # melody and bass come from the tracked lines; alto/tenor by pitch rank below the melody
        rank = {"alto": 1, "tenor": 2}
        for v in pitched:
            r = v["role"]
            i = None
            if r == "melody":
                m = lines["melody"][s]
                i = None if m is None else m - MIDI_LO
            elif r == "bass":
                m = lines["bass"][s]
                i = None if m is None else m - MIDI_LO
                if i is not None and lines["melody"][s] is not None and abs(m - lines["melody"][s]) < 5:
                    i = None
            elif r in rank:
                below = [x for x in by_pitch if lines["melody"][s] is None or x + MIDI_LO < lines["melody"][s]]
                if len(below) > rank[r] - 1:
                    i = below[rank[r] - 1]
                    if lines["bass"][s] is not None and i + MIDI_LO == lines["bass"][s]:
                        i = None
            if i is not None:
                used.add(i)
                emit(v["name"], i, s, col, onset, r == "bass")
        # harmony voices: strongest remaining notes, no octave doubling of used notes
        rest = sorted((i for i in cand if i not in used and all(abs(i - u) != 12 for u in used)), key=lambda i: -col[i])
        for v, i in zip([v for v in pitched if v["role"] == "harmony"], rest):
            emit(v["name"], i, s, col, onset, False)
        prevcol = col

    # drums: onsets of the percussive part per band, on the step grid, fixed pitch
    cache = {}
    for v in voices_spec:
        if v["role"] in DRUMS:
            for st_ in detect_drum(y, v["role"], g, args.drums, cache):
                events[v["name"]].append((st_, NB_LO + DRUMS[v["role"]][1]))
    # clean the melody up (fewer stray notes): key, octave slips, blips, density
    for v in voices_spec:
        if v["role"] == "melody" and args.simplify > 0:
            events[v["name"]] = clean_melody(events[v["name"]], mel_steps, args, spb)
    return events, n_steps, steps


MAJOR = (0, 2, 4, 5, 7, 9, 11)
KEYS = {n.lower(): i for i, n in enumerate(["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"])}
KEYS.update({"db": 1, "eb": 3, "gb": 6, "ab": 8, "bb": 10})
KEYS.update({k + "m": (v + 3) % 12 for k, v in list(KEYS.items())})   # minor key = its relative major scale


def clean_melody(evs: list, conf: np.ndarray, args, spb: int) -> list:
    """post-processing of the melody line:
       1. octave slips: a note an octave away from both neighbours moves back next to them
       2. blips: a one-step note far from both neighbours with weak evidence is dropped
       3. key: weak notes outside the song's scale are dropped
       4. (--simplify 2) at most one note per eighth"""
    if len(evs) < 3:
        return evs
    evs = sorted(evs)
    def strength(s, m):
        i = m - MIDI_LO
        return float(conf[i, s]) if 0 <= i < conf.shape[0] and s < conf.shape[1] else 0.0
    norm = float(np.percentile([strength(s, m) for s, m in evs], 50)) + 1e-9
    out = list(evs)
    for k in range(1, len(out) - 1):
        (_, a), (s1, b), (_, c) = out[k - 1], out[k], out[k + 1]
        for sh in (-12, 12):
            if abs(b - a) >= 10 and abs(b - c) >= 10 and abs(b + sh - a) <= 5 and abs(b + sh - c) <= 5:
                out[k] = (s1, b + sh)
                break
    keep = [True] * len(out)
    for k in range(1, len(out) - 1):
        (_, a), (s1, b), (s2, c) = out[k - 1], out[k], out[k + 1]
        if s2 - s1 <= 1 and abs(b - a) >= 5 and abs(b - c) >= 5 and strength(s1, b) < 1.2 * norm:
            keep[k] = False
    out = [e for e, k in zip(out, keep) if k]
    if args.key != "off":
        pcs = np.bincount([m % 12 for _, m in out], minlength=12).astype(float)
        tonic = max(range(12), key=lambda t: sum(pcs[(t + d) % 12] for d in MAJOR)) if args.key == "auto" else KEYS[args.key.lower()]
        scale = {(tonic + d) % 12 for d in MAJOR}
        out = [(s, m) for s, m in out if m % 12 in scale or strength(s, m) >= 1.5 * norm]
    if args.simplify >= 2 and spb == 4:
        best = {}
        for s, m in out:
            k = s // 2
            if k not in best or strength(s, m) > strength(*best[k]):
                best[k] = (s, m)
        out = sorted(best.values())
    return out


def analyze(y: np.ndarray, sal: np.ndarray, env: np.ndarray, bpm: float, args) -> tuple[str, int, float, list[str]]:
    """decide the voices, steps per beat and note threshold from the recording"""
    notes_ = []
    # steps per beat: are there onsets between the eighth notes?
    g16 = make_grid(sal, env, bpm, 4)
    fps = SR / HOP
    onsets = np.array([env[int(g16["phase"] + s * g16["step_frames"]):int(g16["phase"] + (s + 0.5) * g16["step_frames"]) + 1].max()
                       if int(g16["phase"] + s * g16["step_frames"]) < len(env) else 0 for s in range(g16["n"])])
    eighth, sixteenth = onsets[0::2].mean() + 1e-9, onsets[1::2].mean()
    ratio = sixteenth / eighth
    spb = args.spb or (4 if ratio > 0.45 and bpm * 4 / 60 <= 20 else 2)
    notes_.append(f"off-beat sixteenth onsets {ratio:.2f} of the eighths -> {'1/16' if spb == 4 else '1/8'} grid")
    g = g16 if spb == 4 else make_grid(sal, env, bpm, spb)
    steps, loud = g["steps"], g["loud"]
    # polyphony and bass, on steps that are not silent
    thr = args.threshold
    counts, bass_steps, voiced = [], 0, 0
    for s in range(g["n"]):
        if loud[s] < args.silence:
            continue
        col = steps[:, s] / (steps[:, s].max() + 1e-9)
        c = peaks(col, thr)
        if not c:
            continue
        voiced += 1
        counts.append(len(c))
        if len(c) > 1 and min(c) + MIDI_LO < 60:
            bass_steps += 1
    if voiced == 0:
        return "drums", spb, thr, notes_ + ["no pitched content found"]
    poly = float(np.percentile(counts, 75))
    bass = bass_steps / voiced > 0.3
    # dense mixes: raise the threshold a little, sparse ones: lower it
    if not args.threshold_given:
        thr = 0.62 if poly >= 5 else 0.5 if poly <= 1.5 else thr
    harmony = 1 if poly >= 4 else 0       # more inner voices mostly add mud on note blocks
    notes_.append(f"polyphony ~{poly:.1f} notes, bass on {bass_steps * 100 // voiced}% of steps -> "
                  f"melody{' + bass' if bass else ''}{f' + {harmony} harmony' if harmony else ''}")
    voices = ["melody"] + (["bass"] if bass else []) + ["harmony"] * harmony
    # drums: share of percussive energy and regular hits per bar
    mag = stft_mag(y, 1024)
    perc = percussive(y)
    pshare = float(perc.sum() / (mag.sum() + 1e-9))
    bars = max(1, g["n"] / (spb * 4))
    if pshare > 0.15:
        found = []
        for r in DRUMS:
            hits = len(detect_drum(y, r, g, args.drums, {}))
            need = 3.0 if r == "hat" else 0.75
            if hits / bars >= need:
                voices.append(r)
                found.append(f"{r} {hits / bars:.1f}/bar")
        notes_.append(f"percussive energy {pshare * 100:.0f}% -> " + (", ".join(found) if found else "no regular drums"))
    else:
        notes_.append(f"percussive energy {pshare * 100:.0f}% -> no drums")
    return ",".join(voices), spb, thr, notes_


def detect_drum(y, role, g, threshold, flux_cache) -> list[int]:
    n_steps, phase, step_frames = g["n"], g["phase"], g["step_frames"]
    def drum_steps(sig):
        v = np.zeros(n_steps, np.float32)
        for s in range(n_steps):
            a = int(phase + (s - 0.4) * step_frames)
            b = int(phase + (s + 0.4) * step_frames) + 1
            v[s] = sig[max(0, a):b].max() if max(0, a) < min(b, len(sig)) else 0
        return v / (np.percentile(v, 95) + 1e-9)
    for r in ("kick", "snare", role):
        if r not in flux_cache:
            flux_cache[r] = drum_steps(band_flux(y, DRUMS[r][0]))
    f = flux_cache[role]
    hits, last_hit = [], -9
    for s in range(n_steps):
        left = f[s - 1] if s > 0 else 0
        right = f[s + 1] if s + 1 < n_steps else 0
        if f[s] < threshold or f[s] < left or f[s] <= right:
            continue
        other = {"kick": "snare", "snare": "kick"}.get(role)
        if other and flux_cache[other][s] > f[s] * 1.15 and f[s] < 0.9:
            continue
        if s - last_hit < 2 and f[s] < 1.0:
            continue
        hits.append(s)
        last_hit = s
    return hits


PITCHED_INST = ["harp", "pling", "banjo", "bit", "iron_xylophone", "bass", "didgeridoo", "guitar", "flute", "cow_bell",
                "bell", "chime", "xylophone"]
PLUCKED = {"harp", "pling", "banjo", "bit", "iron_xylophone", "guitar", "bass", "xylophone"}
SUSTAINED = {"flute", "didgeridoo", "bell", "chime"}


def choose_instruments(spec: list, events: dict, steps: np.ndarray) -> list[str]:
    """per voice: the instrument whose real range holds most of the voice's notes
    without an octave shift (so it sounds at the original pitch), with a bonus for
    a matching attack (decaying notes -> plucked, held notes -> sustained) and a
    small penalty for an instrument another voice already uses"""
    used, why = set(), []
    for v in spec:
        if v["inst"] is not None:
            continue
        evs = events.get(v["name"]) or []
        if not evs:
            v["inst"] = "harp"
            continue
        ms = [m for _, m in evs]
        med = float(np.median(ms))
        held = []
        for s, m in evs:
            i = m - MIDI_LO
            if s + 2 < steps.shape[1] and steps[i, s] > 0:
                held.append(min(1.5, steps[i, s + 2] / steps[i, s]))
        sustain = float(np.median(held)) if held else 0.5
        plucked = sustain <= 0.75

        def score(inst):
            base = INST_BASE[inst]
            inside = sum(1 for m in ms if base <= m < base + NB_RANGE) / len(ms)
            bonus = 0.12 if (plucked and inst in PLUCKED) or (not plucked and inst in SUSTAINED) else 0
            if v["role"] == "bass" and inst in ("bass", "didgeridoo", "guitar"):
                bonus += 0.1
            if v["role"] == "melody" and inst in ("harp", "flute", "pling", "bell", "xylophone"):
                bonus += 0.05
            return inside + bonus - (0.2 if inst in used else 0)
        pick = max(PITCHED_INST, key=score)
        base = INST_BASE[pick]
        inside = sum(1 for m in ms if base <= m < base + NB_RANGE) * 100 // len(ms)
        used.add(pick)
        v["inst"] = pick
        why.append(f"{v['name']}: median {NAMES[int(med) % 12]}{int(med) // 12 - 1}, {'plucked' if plucked else 'held'} "
                   f"-> {pick} ({inside}% at the original pitch)")
    return why


NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]


def build_song(events: dict, voices_spec: list, n_steps: int, bpm: float, spb: int, args) -> dict:
    tracks, notes, stats = [], {}, {}
    ti = 0
    for v in voices_spec:
        evs = events[v["name"]]
        if not evs and not args.keep_empty:
            stats[v["name"]] = (0, 0, 0)
            continue
        ti += 1
        drum = v["role"] in DRUMS
        base = NB_LO if drum else INST_BASE.get(v["inst"], NB_LO)
        shift = (0 if drum else fold([m for _, m in evs], base) + (args.transpose or 0)) + v["shift"]
        inside = 0
        for s, m in evs:
            mm = m + shift
            if base <= mm < base + NB_RANGE:
                inside += 1
            step = s + 1
            lst = notes.setdefault(step, [])
            if any(t == ti for t, _ in lst):          # one note per track per step
                continue
            lst.append((ti, to_block(mm, base) if not drum else to_block(mm)))
        tracks.append({"inst": v["inst"], "name": v["name"]})
        stats[v["name"]] = (len(evs), shift, inside)
    bar = spb * 4
    last = max(notes) if notes else bar
    length = ((last - 1) // bar + 1) * bar
    return {"bpm": int(round(bpm)), "spb": spb, "length": length, "tracks": tracks, "notes": notes, "stats": stats}


def lua_song(song: dict) -> str:
    def q(s: str) -> str:
        return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"'
    tracks = ",".join("{inst=%s,name=%s}" % (q(t["inst"]), q(t["name"])) for t in song["tracks"])
    notes = ",".join("[%d]={%s}" % (s, ",".join("{%d,%d}" % n for n in song["notes"][s])) for s in sorted(song["notes"]))
    return "{bpm=%d,spb=%d,length=%d,tracks={%s},notes={%s}}" % (song["bpm"], song["spb"], song["length"], tracks, notes)


# ---------------------------------------------------------------- preview
def render_preview(song: dict, path: str) -> None:
    sr = 22050
    step = 60 / song["bpm"] / song["spb"]
    total = int((song["length"] + 4) * step * sr)
    out = np.zeros(total, np.float32)
    t = np.arange(int(0.6 * sr)) / sr
    rng = np.random.default_rng(1)
    noise = rng.uniform(-1, 1, len(t)).astype(np.float32)
    for s, lst in song["notes"].items():
        for ti, p in lst:
            inst = song["tracks"][ti - 1]["inst"]
            f = 440.0 * 2 ** ((INST_BASE.get(inst, NB_LO) + p - 69) / 12)
            if inst == "basedrum":
                tone = np.sin(2 * np.pi * 60 * t * np.exp(-t * 8)) * np.exp(-t * 18)
            elif inst == "snare":
                tone = (0.7 * noise + 0.3 * np.sin(2 * np.pi * 190 * t)) * np.exp(-t * 25) * 0.6
            elif inst == "hat":
                tone = np.diff(noise, prepend=0) * np.exp(-t * 60) * 0.3
            else:
                tone = np.sin(2 * np.pi * f * t) + 0.3 * np.sin(4 * np.pi * f * t)
                tone *= np.exp(-t * (6 if ti == 1 else 4)) * (0.5 if ti > 1 else 0.6)
            a = int((s - 1) * step * sr)
            b = min(total, a + len(tone))
            out[a:b] += tone[:b - a]
    out /= max(1e-9, np.max(np.abs(out)))
    pcm = (out * 32000).astype("<i2")
    import wave
    with wave.open(path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(pcm.tobytes())


# ---------------------------------------------------------------- main
# ================================================================ MIDI input
# A MIDI file already contains exact notes: no pitch detection, only mapping to
# note blocks: tracks/channels -> voices, chords -> several voices, GM programs ->
# note block instruments, drums (channel 10) -> kick/snare/hat, timing -> grid.

def read_midi(path: str) -> dict:
    """minimal Standard MIDI File reader (format 0/1). Returns
    {ppq, tempo (us per quarter, first), tracks: [{name, notes: [(tick, ch, note, vel)], programs: {ch: prog}}]}"""
    data = open(path, "rb").read()
    if data[:4] != b"MThd":
        sys.exit("not a MIDI file")
    hlen = int.from_bytes(data[4:8], "big")
    fmt, ntrk, div = int.from_bytes(data[8:10], "big"), int.from_bytes(data[10:12], "big"), int.from_bytes(data[12:14], "big")
    if div & 0x8000:
        sys.exit("SMPTE-timed MIDI files are not supported")
    pos = 8 + hlen
    tracks, tempo, tempos = [], None, []
    for _ in range(ntrk):
        while data[pos:pos + 4] != b"MTrk":
            if pos >= len(data):
                break
            pos += 8 + int.from_bytes(data[pos + 4:pos + 8], "big")      # skip unknown chunks
        if pos >= len(data):
            break
        ln = int.from_bytes(data[pos + 4:pos + 8], "big")
        p, end = pos + 8, pos + 8 + ln
        pos = end
        tick, status = 0, 0
        tr = {"name": "", "notes": [], "programs": {}}
        def vlq():
            nonlocal p
            v = 0
            while True:
                b = data[p]; p += 1
                v = (v << 7) | (b & 0x7F)
                if not b & 0x80:
                    return v
        while p < end:
            tick += vlq()
            b = data[p]
            if b & 0x80:
                status = b; p += 1
            if status == 0xFF:
                typ = data[p]; p += 1
                n = vlq()
                body = data[p:p + n]; p += n
                if typ == 0x03 and not tr["name"]:
                    tr["name"] = body.decode("latin-1", "replace").strip()
                elif typ == 0x51 and n == 3:
                    us = int.from_bytes(body, "big")
                    tempos.append((tick, us))
                    if tempo is None:
                        tempo = us
                elif typ == 0x2F:
                    break
                status = 0
            elif status in (0xF0, 0xF7):
                n = vlq(); p += n; status = 0
            else:
                kind, ch = status & 0xF0, status & 0x0F
                if kind in (0xC0, 0xD0):
                    d1 = data[p]; p += 1
                    if kind == 0xC0:
                        tr["programs"].setdefault(ch, d1)
                else:
                    d1, d2 = data[p], data[p + 1]; p += 2
                    if kind == 0x90 and d2 > 0:
                        tr["notes"].append((tick, ch, d1, d2))
        tracks.append(tr)
    return {"ppq": div, "tempo": tempo or 500000, "tempos": tempos, "tracks": tracks, "format": fmt}


# General MIDI program families -> note block instruments that suit them
GM_FAMILY = [
    (0, 7, ["harp", "pling", "banjo"]),                  # pianos
    (8, 8, ["bell", "chime"]), (9, 9, ["bell", "chime"]), (10, 10, ["chime", "bell"]),     # celesta, glockenspiel, music box
    (11, 11, ["iron_xylophone"]), (12, 12, ["iron_xylophone", "xylophone"]), (13, 13, ["xylophone"]),
    (14, 15, ["chime", "bell"]),                         # tubular bells, dulcimer
    (16, 23, ["flute", "bit", "didgeridoo"]),            # organs, accordion, harmonica
    (24, 31, ["guitar", "harp", "banjo"]),               # guitars
    (32, 39, ["bass", "didgeridoo", "guitar"]),          # basses
    (40, 55, ["flute", "harp", "didgeridoo", "bell"]),   # strings, ensembles, choir
    (56, 63, ["didgeridoo", "bit", "flute"]),            # brass
    (64, 79, ["flute", "didgeridoo"]),                   # reeds, pipes
    (80, 87, ["bit", "pling", "flute"]),                 # synth leads
    (88, 103, ["pling", "chime", "bell", "flute"]),      # pads, synth fx
    (104, 104, ["banjo"]), (105, 105, ["banjo"]), (106, 111, ["banjo", "harp", "didgeridoo"]),  # ethnic
    (112, 112, ["bell"]), (113, 113, ["cow_bell"]), (114, 119, ["xylophone", "basedrum", "cow_bell"]),
    (120, 127, ["pling", "hat"]),
]
GM_DRUM = {35: "kick", 36: "kick", 41: "kick", 43: "kick", 45: "kick", 47: "kick", 48: "kick", 50: "kick",
           37: "snare", 38: "snare", 39: "snare", 40: "snare",
           42: "hat", 44: "hat", 46: "hat", 49: "hat", 51: "hat", 52: "hat", 53: "hat", 55: "hat", 57: "hat", 59: "hat"}
DRUM_PITCH = {"kick": 3, "snare": 10, "hat": 18}


def midi_instrument(prog: int | None, notes: list[int], used: set) -> str:
    cands = next((c for lo, hi, c in GM_FAMILY if prog is not None and lo <= prog <= hi), ["harp", "flute", "bell", "guitar", "bass"])
    def score(inst):
        base = INST_BASE[inst]
        inside = sum(1 for m in notes if base <= m < base + NB_RANGE) / max(1, len(notes))
        return inside - (0.15 if inst in used else 0) - 0.02 * cands.index(inst)
    return max(cands, key=score)


def convert_midi(args, out: str) -> None:
    mf = read_midi(args.input)
    ppq = mf["ppq"]
    bpm = args.bpm or round(60_000_000 / mf["tempo"], 2)
    if len({us for _, us in mf["tempos"]}) > 1:
        print(f"note     : the file changes tempo {len(mf['tempos'])} times; Composer plays one tempo ({bpm:g} bpm)")
    # split into parts: one per (track, channel); channel 10 (index 9) = drums
    parts = []
    for ti, tr in enumerate(mf["tracks"], start=1):
        chans = sorted({ch for _, ch, _, _ in tr["notes"]})
        for ch in chans:
            notes = [(t, n, v) for t, c, n, v in tr["notes"] if c == ch]
            prog = tr["programs"].get(ch)
            if prog is None:
                for other in mf["tracks"]:
                    if ch in other["programs"]:
                        prog = other["programs"][ch]
                        break
            name = tr["name"] or f"track {ti}"
            if len(chans) > 1:
                name += f" ch{ch + 1}"
            parts.append({"index": len(parts) + 1, "name": name, "ch": ch, "prog": prog, "notes": notes, "drums": ch == 9})
    if not parts:
        sys.exit("the MIDI file has no notes")
    # start/duration (seconds at the main tempo)
    tick_s = 60 / bpm / ppq
    t0 = int(args.start / tick_s)
    t1 = int((args.start + args.duration) / tick_s) if args.duration else None
    # grid: the coarsest steps per beat that keeps (almost) all onsets exact
    onsets = [t for p in parts for t, _, _ in p["notes"] if t >= t0 and (t1 is None or t < t1)]
    spb = args.spb
    if not spb:
        for cand in (1, 2, 3, 4, 6, 8):
            step = ppq / cand
            exact = sum(1 for t in onsets if abs(t / step - round(t / step)) < 0.1) / max(1, len(onsets))
            if exact >= 0.97 or cand == 8:
                spb = cand
                break
        while bpm * spb / 60 > 20 and spb > 1:      # note blocks: max 20 notes per second per block
            spb = {8: 4, 6: 3, 4: 2, 3: 2, 2: 1}[spb]
        spb = max(spb, 2) if bpm * 2 / 60 <= 20 else spb
    step_ticks = ppq / spb
    print(f"MIDI     : {len(parts)} parts, {ppq} ticks per beat, {bpm:g} bpm, grid 1/{spb * 4} ({bpm * spb / 60:.1f} steps/s)")
    # choose parts
    sel = parts
    if args.parts:
        want = set()
        for x in args.parts.split(","):
            x = x.strip()
            want |= {p["index"] for p in parts if x.isdigit() and p["index"] == int(x) or (not x.isdigit() and x.lower() in p["name"].lower())}
        sel = [p for p in parts if p["index"] in want]
    if args.no_drums:
        sel = [p for p in sel if not p["drums"]]
    for p in parts:
        mark = "*" if p in sel else " "
        kind = "drums" if p["drums"] else f"program {p['prog'] if p['prog'] is not None else '-'}"
        rng = [n for _, n, _ in p["notes"]]
        print(f"  {mark}{p['index']:2}. {p['name'][:28]:28} {kind:12} {len(p['notes']):5} notes  {min(rng)}..{max(rng)}")
    # voices
    tracks, notes, used = [], {}, set()
    def add_note(ti, step, pitch):
        lst = notes.setdefault(step, [])
        if not any(t == ti for t, _ in lst):
            lst.append((ti, pitch))
    overrides = dict(x.split("=", 1) for x in args.inst.split(",")) if args.inst else {}
    for p in sel:
        evs = [(int(round((t - t0) / step_ticks)) + 1, n, v) for t, n, v in p["notes"] if t >= t0 and (t1 is None or t < t1)]
        if args.min_velocity:
            evs = [e for e in evs if e[2] >= args.min_velocity]
        if not evs:
            continue
        if p["drums"]:
            for kind in ("kick", "snare", "hat"):
                ds = [(s, n) for s, n, _ in evs if GM_DRUM.get(n) == kind]
                if not ds:
                    continue
                tracks.append({"inst": {"kick": "basedrum", "snare": "snare", "hat": "hat"}[kind], "name": f"{p['name']} {kind}"})
                ti = len(tracks)
                for s, _ in ds:
                    add_note(ti, s, DRUM_PITCH[kind])
                print(f"  -> {tracks[-1]['name'][:30]:30} {tracks[-1]['inst']:14} {len(ds):5} hits")
            continue
        # chords: per step the notes from high to low go to voice 1, 2, 3…; extra notes
        # beyond --max-poly are dropped from the middle (keep top and bottom)
        by_step = {}
        for s, n, v in evs:
            by_step.setdefault(s, set()).add(n)
        poly = min(args.max_poly, max(len(v) for v in by_step.values()))
        voices = [[] for _ in range(poly)]
        for s, ns in by_step.items():
            ns = sorted(ns, reverse=True)
            if len(ns) > poly:
                ns = ns[:poly - 1] + [ns[-1]] if poly > 1 else ns[:1]
            for k, n in enumerate(ns):
                voices[k].append((s, n))
        for k, vs in enumerate(voices):
            if not vs:
                continue
            name = p["name"] if poly == 1 else f"{p['name']} {k + 1}"
            inst = (overrides.get(name) or overrides.get(str(p["index"])) or overrides.get(p["name"])
                    or midi_instrument(p["prog"], [n for _, n in vs], used))
            if inst not in INSTRUMENTS:
                sys.exit(f"unknown instrument {inst!r}")
            used.add(inst)
            base = INST_BASE[inst]
            shift = fold([n for _, n in vs], base) + args.transpose
            inside = sum(1 for _, n in vs if base <= n + shift < base + NB_RANGE) * 100 // len(vs)
            tracks.append({"inst": inst, "name": name})
            ti = len(tracks)
            for s, n in vs:
                add_note(ti, s, to_block(n + shift, base))
            print(f"  -> {name[:30]:30} {inst:14} {len(vs):5} notes, shift {shift:+d}, {inside}% at that octave")
    if not tracks:
        sys.exit("nothing to convert (check --parts / --start)")
    bar = spb * 4
    last = max(notes)
    song = {"bpm": int(round(bpm)), "spb": spb, "length": ((last - 1) // bar + 1) * bar, "tracks": tracks, "notes": notes}
    with open(out, "w", encoding="utf-8") as f:
        f.write(lua_song(song))
    print(f"Written  : {out}  ({song['length']} steps, {sum(len(v) for v in notes.values())} notes, {len(tracks)} tracks)")
    if args.preview:
        render_preview(song, args.preview)
        print(f"Preview  : {args.preview}")


def main() -> None:
    ap = argparse.ArgumentParser(description="audio (mp3/wav/...) -> .yms song for the Yaru Composer (note blocks)")
    ap.add_argument("input")
    ap.add_argument("output", nargs="?", help="default: <input>.yms")
    ap.add_argument("--start", type=float, default=0.0, help="start offset, seconds")
    ap.add_argument("--duration", type=float, help="length to convert, seconds")
    ap.add_argument("--bpm", type=float, help="tempo (default: detected)")
    ap.add_argument("--spb", type=int, default=None, choices=(2, 4), help="steps per beat: 2 = eighths, 4 = sixteenths (default: auto)")
    ap.add_argument("--voices", default="auto", help="'auto' (default: decided from the audio) or voices = tracks, comma separated: role[:instrument[:semitones]]. "
                    "roles: melody alto tenor bass harmony kick snare hat (repeat harmony for more chord notes). "
                    "presets: " + ", ".join(f"{k}={v}" for k, v in PRESETS.items()))
    ap.add_argument("--tracks", type=int, choices=(1, 2, 3), help="old shortcut: 1 melody, 2 +bass, 3 +harmony")
    ap.add_argument("--drums", type=float, default=0.55, help="drum hit threshold 0..1 (higher = fewer hits)")
    ap.add_argument("--keep-empty", action="store_true", help="keep voices that got no notes as empty tracks")
    ap.add_argument("--threshold", type=float, default=None, help="relative note strength 0..1 (higher = fewer notes; default: auto around 0.55)")
    ap.add_argument("--silence", type=float, default=0.12, help="steps quieter than this are rests")
    ap.add_argument("--melody-low", dest="mel_lo", type=int, default=55, help="lowest MIDI note the melody may use (default 55 = G3)")
    ap.add_argument("--onset", type=float, default=0.6, help="onset strength that re-strikes held notes")
    ap.add_argument("--restrike", type=int, default=0, help="re-strike held notes every N steps (0 = off)")
    ap.add_argument("--transpose", type=int, default=0, help="extra semitones after octave folding")
    ap.add_argument("--preview", help="also write a WAV preview of the arrangement")
    ap.add_argument("--simplify", type=int, default=1, choices=(0, 1, 2),
                    help="melody clean-up: 0 off, 1 fix octave slips/blips/out-of-key notes (default), 2 also max one note per eighth")
    ap.add_argument("--key", default="auto", help="song key for the clean-up: auto, off, or e.g. C, Am, F#, Bbm")
    ap.add_argument("--parts", help="MIDI: which parts to keep, by number or name, e.g. 1,3 or piano,bass")
    ap.add_argument("--no-drums", action="store_true", help="MIDI: leave the drum channel out")
    ap.add_argument("--max-poly", type=int, default=3, help="MIDI: max simultaneous notes per part (each becomes a track)")
    ap.add_argument("--inst", help="MIDI: instrument per part, e.g. 1=flute,bass=didgeridoo")
    ap.add_argument("--min-velocity", type=int, default=0, help="MIDI: ignore quieter notes (1..127)")
    ap.add_argument("--separate", choices=("off", "vocals", "other"), default="off",
                    help="isolate the melody first with Demucs (py -m pip install demucs): vocals = sung melody, other = lead instruments")
    ap.add_argument("--engine", choices=("auto", "nn", "dsp"), default="auto",
                    help="pitch detection: nn = Basic Pitch neural network (needs onnxruntime + model), dsp = built-in; auto = nn if available")
    ap.add_argument("--model", help="path to Basic Pitch nmp.onnx (default: next to this script or from the basic-pitch package)")
    args = ap.parse_args()

    if not os.path.isfile(args.input):
        folder = os.path.dirname(os.path.abspath(args.input))
        found = sorted(f for f in os.listdir(folder) if f.lower().endswith((".mid", ".midi", ".kar", ".mp3", ".wav", ".ogg", ".flac"))) \
            if os.path.isdir(folder) else []
        msg = f"file not found: {args.input}"
        if found:
            msg += "\nmusic files in " + folder + ":\n  " + "\n  ".join(found[:20])
        sys.exit(msg)
    out = args.output or os.path.splitext(args.input)[0] + ".yms"
    if args.input.lower().endswith((".mid", ".midi", ".kar")):
        return convert_midi(args, out)
    y = load_audio(args.input, args.start, args.duration)
    if len(y) < SR:
        sys.exit("audio is shorter than one second")
    print(f"Audio    : {len(y) / SR:.1f} s")
    mag = stft_mag(y)                      # long window: pitch
    env = onset_envelope(stft_mag(y, 1024))  # short window: timing
    bpm = args.bpm or estimate_tempo(env)
    sal = salience(mag)
    args.onset_post = None
    args.melody_audio = None
    if args.separate != "off":
        args.melody_audio = load_audio(separate(args.input, args.separate), args.start, args.duration)[:len(y)]
        if len(args.melody_audio) < len(y):
            args.melody_audio = np.pad(args.melody_audio, (0, len(y) - len(args.melody_audio)))
        print(f"Separate : melody taken from the '{args.separate}' stem")
    model = find_model(args.model) if args.engine in ("auto", "nn") else None
    if args.engine == "nn" and not model:
        sys.exit("neural engine: model not found (py -m pip install basic-pitch, or --model path/to/nmp.onnx)")
    if model:
        try:
            note_post, onset_post = nn_posteriors(y, model, sal.shape[1])
            sal, args.onset_post = note_post, onset_post
            if args.melody_audio is not None:
                args.mel_note, args.onset_post = nn_posteriors(args.melody_audio, model, sal.shape[1])
            print(f"Engine   : neural (Basic Pitch, {os.path.basename(model)})")
        except ImportError:
            if args.engine == "nn":
                sys.exit("neural engine needs onnxruntime: py -m pip install onnxruntime")
            print("Engine   : signal processing (install onnxruntime + basic-pitch for the neural engine)")
    else:
        print("Engine   : signal processing (py -m pip install basic-pitch onnxruntime for the more accurate neural engine)")
    args.threshold_given = args.threshold is not None
    if args.threshold is None:
        args.threshold = 0.55
    auto = args.voices == "auto" and not args.tracks
    if auto:
        voices, spb, args.threshold, why = analyze(y, sal, env, bpm, args)
        for line in why:
            print(f"Analysis : {line}")
    else:
        voices, spb = args.voices if args.voices != "auto" else str(args.tracks), args.spb or 4
    args.spb = spb
    rate = bpm * spb / 60
    if rate > 20:
        print(f"warning: {rate:.1f} steps/s is above 20 (one note block note per tick); use --spb 2 or a lower --bpm")
    print(f"Tempo    : {bpm:.1f} bpm{' (detected)' if not args.bpm else ''}, {spb} steps per beat ({rate:.1f} steps/s), threshold {args.threshold:.2f}")
    spec = parse_voices(voices)
    events, n_steps, steps = transcribe(sal, env, bpm, spb, args, spec, y)
    for line in choose_instruments(spec, events, steps):
        print(f"Sound    : {line}")
    print("Voices   : " + ", ".join(f"{v['name']} ({v['inst']})" for v in spec))
    song = build_song(events, spec, n_steps, bpm, spb, args)
    for voice, (n, shift, inside) in song["stats"].items():
        if n == 0:
            print(f"{voice:10}: no notes found (track left out; --keep-empty keeps it)")
        else:
            print(f"{voice:10}: {n} notes, octave shift {shift:+d}, {inside * 100 // max(1, n)}% in the instrument's range")
    with open(out, "w", encoding="utf-8") as f:
        f.write(lua_song(song))
    print(f"Written  : {out}  ({song['length']} steps, {sum(len(v) for v in song['notes'].values())} notes)")
    if args.preview:
        render_preview(song, args.preview)
        print(f"Preview  : {args.preview}")
    print("Copy the .yms to /home/Music on the OC computer and open it in Composer.")


if __name__ == "__main__":
    main()