#!/usr/bin/env python3
"""
pianoify — turn a song recording (MP3/WAV/…) into a simplified piano arrangement.

Pipeline
  1. Load + separate      harmonic/percussive split (optionally Demucs stems if installed)
  2. Rhythm               beat tracking, tempo, downbeat HMM that re-syncs bar lines
  3. Key                  Krumhansl-Schmuckler on the whole-song chroma
  4. Chords               beat-synchronous chroma + bass chroma → template scores
                          → HMM (Viterbi) smoothing with a diatonic prior
  5. Melody               predominant-pitch tracking (harmonic salience + Viterbi)
                          → quantised to an 8th (standard) or quarter (easy) grid
  6. Arrange              RH = melody (folded into a comfortable range),
                          LH = voice-led triads (standard) or root+fifth (easy)
  7. Engrave              music21 → MusicXML + MIDI, Verovio → PDF/SVG sheet music,
                          plus a plain-text chord chart and a JSON analysis

Usage
  python pianoify.py song.mp3 -o out/            # standard arrangement
  python pianoify.py song.mp3 --level easy       # beginner: quarter-note grid, root+5th LH
  python pianoify.py song.mp3 --to-key C         # transpose to an easier key
  python pianoify.py song.mp3 --transpose -2     # or by a number of semitones
  python pianoify.py song.mp3 --bpm 92 --meter 3 # override tempo / beats per bar
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
import warnings
from collections import Counter
from dataclasses import asdict, dataclass, field

import numpy as np

warnings.filterwarnings("ignore")

import librosa  # noqa: E402

SR = 22050
HOP = 512

SHARP = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
FLAT = ["C", "Db", "D", "Eb", "E", "F", "Gb", "G", "Ab", "A", "Bb", "B"]
FLAT_KEYS = {"F", "Bb", "Eb", "Ab", "Db", "Gb", "Dm", "Gm", "Cm", "Fm", "Bbm", "Ebm"}

# chord qualities: intervals + per-tone template weights
QUALITIES = {
    "": ([0, 4, 7], [1.0, 0.8, 0.9]),
    "m": ([0, 3, 7], [1.0, 0.8, 0.9]),
    "7": ([0, 4, 7, 10], [1.0, 0.8, 0.8, 0.6]),
    "maj7": ([0, 4, 7, 11], [1.0, 0.8, 0.8, 0.6]),
    "m7": ([0, 3, 7, 10], [1.0, 0.8, 0.8, 0.6]),
}

# Krumhansl-Kessler key profiles
KK_MAJOR = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88])
KK_MINOR = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])
MAJOR_SCALE = [0, 2, 4, 5, 7, 9, 11]
MINOR_SCALE = [0, 2, 3, 5, 7, 8, 10]


# ──────────────────────────────────────────────────────────────── data types
@dataclass
class Chord:
    root: int | None          # pitch class, None = no chord
    quality: str = ""
    start_beat: float = 0.0   # in beats from the start of bar 1 (may be < 0 for pickup)
    beats: float = 0.0

    def name(self, flats=False) -> str:
        if self.root is None:
            return "N.C."
        return (FLAT if flats else SHARP)[self.root] + self.quality


@dataclass
class Note:
    midi: int
    start_beat: float
    beats: float


@dataclass
class Analysis:
    title: str
    duration_sec: float
    bpm: float
    beats_per_bar: int
    key_tonic: int
    key_mode: str               # "major" | "minor"
    transpose: int = 0
    chords: list = field(default_factory=list)
    melody: list = field(default_factory=list)
    beat_times: list = field(default_factory=list)
    origin_beat: int = 0        # index into beat_times that is beat 0 of bar 1

    @property
    def key_name(self):
        n = (FLAT if self.uses_flats else SHARP)[self.key_tonic]
        return f"{n} {self.key_mode}"

    @property
    def uses_flats(self):
        n = FLAT[self.key_tonic] + ("m" if self.key_mode == "minor" else "")
        return n in FLAT_KEYS


# ──────────────────────────────────────────────────────────────── 1. audio
def load_audio(path, separate=False):
    y, _ = librosa.load(path, sr=SR, mono=True)
    stems = {}
    if separate:
        stems = try_demucs(path)
    y_harm, y_perc = librosa.effects.hpss(y, margin=(1.0, 2.0))
    return y, y_harm, y_perc, stems


def try_demucs(path):
    """Optional: use Demucs (if installed) to get cleaner vocal/accompaniment stems."""
    try:
        import torch  # noqa: F401
        from demucs.api import Separator  # demucs >= 4.1
    except Exception:
        print("  (demucs not installed — continuing with HPSS only)", file=sys.stderr)
        return {}
    try:
        sep = Separator(model="htdemucs")
        _, stems = sep.separate_audio_file(path)
    except Exception as e:  # e.g. model weights can't be downloaded
        print(f"  (demucs failed: {e.__class__.__name__}: {str(e)[:80]} — using HPSS only)", file=sys.stderr)
        return {}
    out = {}
    for name, wav in stems.items():
        mono = wav.mean(0).cpu().numpy()
        out[name] = librosa.resample(mono, orig_sr=sep.samplerate, target_sr=SR)
    return out


# ──────────────────────────────────────────────────────────────── 2. rhythm
def track_beats(y, y_perc, bpm=None):
    onset_env = librosa.onset.onset_strength(y=y_perc, sr=SR, hop_length=HOP, aggregate=np.median)
    kw = {"start_bpm": bpm, "tightness": 400} if bpm else {}
    tempo, frames = librosa.beat.beat_track(onset_envelope=onset_env, sr=SR, hop_length=HOP,
                                            trim=False, **kw)
    times = librosa.frames_to_time(frames, sr=SR, hop_length=HOP)
    tempo = float(np.atleast_1d(tempo)[0])
    if len(times) < 4:  # degenerate: fall back to a fixed grid
        period = 60.0 / (bpm or tempo or 100)
        times = np.arange(0, len(y) / SR, period)
    period = float(np.median(np.diff(times)))
    # extend the grid to cover the whole file so nothing falls outside a beat
    dur = len(y) / SR
    pre = np.arange(times[0] - period, -1e-6, -period)[::-1]
    post = np.arange(times[-1] + period, dur + period, period)
    times = np.concatenate([pre, times, post])
    return times, 60.0 / period, onset_env


def per_beat(feature, beat_times, agg=np.mean):
    """Aggregate a (d, frames) feature into one column per beat interval."""
    fr = librosa.time_to_frames(beat_times, sr=SR, hop_length=HOP)
    fr = np.clip(fr, 0, feature.shape[1])
    cols = []
    for a, b in zip(fr[:-1], fr[1:]):
        cols.append(agg(feature[:, a:max(b, a + 1)], axis=1) if a < feature.shape[1]
                    else np.zeros(feature.shape[0]))
    return np.array(cols).T  # (d, n_beats-1)


# ──────────────────────────────────────────────────────────────── 3. key
def estimate_key(chroma_sum):
    best = (-2, 0, "major")
    c = chroma_sum - chroma_sum.mean()
    for t in range(12):
        for mode, prof in (("major", KK_MAJOR), ("minor", KK_MINOR)):
            p = np.roll(prof, t)
            r = np.corrcoef(c, p - p.mean())[0, 1]
            if r > best[0]:
                best = (r, t, mode)
    return best[1], best[2]


def diatonic_chords(tonic, mode):
    """Set of (root, quality) that are diatonic triads / common 7ths in the key."""
    if mode == "major":
        degs = [(0, ""), (2, "m"), (4, "m"), (5, ""), (7, ""), (9, "m"),
                (7, "7"), (0, "maj7"), (5, "maj7"), (2, "m7"), (9, "m7"), (4, "m7")]
    else:
        degs = [(0, "m"), (3, ""), (5, "m"), (7, "m"), (7, ""), (8, ""), (10, ""),
                (7, "7"), (0, "m7"), (5, "m7"), (3, "maj7"), (8, "maj7")]
    return {((tonic + d) % 12, q) for d, q in degs}


# ──────────────────────────────────────────────────────────────── 4. chords
def chord_vocab(sevenths):
    vocab = [(None, "N")]
    for q in (QUALITIES if sevenths else ["", "m"]):
        for r in range(12):
            vocab.append((r, q))
    return vocab


def chord_templates(vocab):
    T = np.zeros((len(vocab), 12))
    for i, (r, q) in enumerate(vocab):
        if r is None:
            T[i] = 1.0
            continue
        ivs, w = QUALITIES[q]
        for iv, wt in zip(ivs, w):
            T[i, (r + iv) % 12] = wt
    return T / np.linalg.norm(T, axis=1, keepdims=True)


def detect_chords(chroma_b, bass_b, energy_b, key, sevenths=False, p_stay=0.85):
    vocab = chord_vocab(sevenths)
    T = chord_templates(vocab)
    C = chroma_b / (np.linalg.norm(chroma_b, axis=0, keepdims=True) + 1e-9)
    scores = T @ C                                          # (n_chords, n_beats)

    # bass evidence: does the lowest register agree with the chord root?
    B = bass_b / (bass_b.sum(axis=0, keepdims=True) + 1e-9)
    diat = diatonic_chords(*key)
    for i, (r, q) in enumerate(vocab):
        if r is None:
            continue
        scores[i] += 0.25 * B[r]
        if (r, q) in diat:
            scores[i] += 0.04                               # gentle key prior
        if q not in ("", "m"):
            scores[i] -= 0.03                               # prefer plain triads

    # silence / non-harmonic beats → N.C.
    quiet = energy_b < 0.08 * np.percentile(energy_b, 90)
    raw = T[1:] @ C
    atonal = (raw.max(axis=0) - np.median(raw, axis=0)) < 0.22         # flat chroma = drums/noise
    scores[0] = np.where(quiet | atonal, 2.0, 0.35)

    probs = np.exp(25 * (scores - scores.max(axis=0, keepdims=True)))
    probs /= probs.sum(axis=0, keepdims=True)
    trans = librosa.sequence.transition_loop(len(vocab), p_stay)
    path = librosa.sequence.viterbi(probs, trans)
    return [vocab[i] for i in path]


def track_downbeats(labels, chroma_b, low_onset_b, beats_per_bar, p_resync=0.004):
    """HMM over position-in-bar. Beat 1 is where harmony changes and the bass hits;
    the model normally advances one position per beat but may re-sync, so a dropped
    beat or an odd-length bar doesn't shift every bar line after it."""
    n, B = len(labels), beats_per_bar
    C = chroma_b / (np.linalg.norm(chroma_b, axis=0, keepdims=True) + 1e-9)
    flux = np.r_[0, 1 - np.sum(C[:, 1:] * C[:, :-1], axis=0)]
    flux = flux / (np.percentile(flux, 95) + 1e-9)
    change = np.array([0.0] + [float(labels[i] != labels[i - 1]) for i in range(1, n)])
    lo = low_onset_b / (np.percentile(low_onset_b, 95) + 1e-9)
    ev = 2.0 * change + np.clip(flux, 0, 2) + 0.7 * np.clip(lo, 0, 2)
    logp = np.zeros((B, n))
    logp[0] = 1.2 * ev
    if B % 2 == 0:
        logp[B // 2] = 0.6 * ev          # half-bar changes are common but less likely
    P = np.exp(logp - logp.max(axis=0, keepdims=True))
    P /= P.sum(axis=0, keepdims=True)
    trans = np.full((B, B), p_resync / (B - 1))
    for s in range(B):
        trans[s, (s + 1) % B] = 1 - p_resync
    pos = librosa.sequence.viterbi(P, trans, p_init=np.full(B, 1 / B))
    return [i for i in range(n) if pos[i] == 0]


def regrid(beat_times, downbeats, B):
    """New beat grid with exactly B evenly spaced beats between consecutive downbeats.
    Returns (grid, index of the first downbeat in the grid)."""
    if len(downbeats) < 2:
        return np.asarray(beat_times), 0
    period = float(np.median(np.diff(beat_times)))
    d0 = beat_times[downbeats[0]]
    pre = list(np.arange(d0 - period, -1e-6, -period))[::-1]
    out = pre[:]
    for a, b in zip(downbeats[:-1], downbeats[1:]):
        out.extend(np.linspace(beat_times[a], beat_times[b], B, endpoint=False))
    t = beat_times[downbeats[-1]]
    while t <= beat_times[-1] + 1e-6:
        out.append(t)
        t += period
    return np.array(out), len(pre)


def segment_chords(labels, origin, beats_per_bar, resolution):
    """Collapse per-beat labels to one chord per `resolution` beats, then merge repeats."""
    n = len(labels)
    first = origin - int(np.ceil(origin / beats_per_bar)) * beats_per_bar if origin else 0
    out = []
    for s in range(first, n, resolution):
        win = [labels[i] for i in range(max(s, 0), min(s + resolution, n))]
        if not win:
            continue
        cnt = Counter(win)
        for k in cnt:                                       # a chord beats "no chord" on ties
            if k[0] is None:
                cnt[k] *= 0.49
        top = max(cnt.values())
        lab = next(l for l in win if cnt[l] == top)         # ties → earliest in window
        out.append(Chord(lab[0], "" if lab[1] == "N" else lab[1], float(s - origin), float(resolution)))
    merged = []
    for c in out:
        if merged and (merged[-1].root, merged[-1].quality) == (c.root, c.quality):
            merged[-1].beats += c.beats
        else:
            merged.append(c)
    # drop leading/trailing N.C.
    while merged and merged[0].root is None:
        merged.pop(0)
    while merged and merged[-1].root is None:
        merged.pop()
    return merged


# ──────────────────────────────────────────────────────────────── 5. melody
MEL_BPO = 36                 # 1/3-semitone resolution
MEL_FMIN_MIDI = 36           # C2: lowest CQT bin
MEL_LO, MEL_HI = 48, 84      # melody f0 search range C3..C6
MEL_MEDIAN_SEC = 4.0         # accompaniment-suppression window
MEL_MEDIAN_AMT = 0.6


def vibrato_score(S):
    """0..1 per (pitch candidate, frame): how much this pitch is being sung with vibrato.

    Vibrato is frequency modulation (4-8 Hz, about +-30 cents), so energy swings between the
    candidate bins just above and just below in opposite phase. Restruck or tremolo notes are
    amplitude modulation: neighbours rise and fall together. So: band-pass each candidate's
    salience to 4-8 Hz, and score anti-phase correlation between its neighbours.
    """
    from scipy.signal import butter, sosfiltfilt
    from scipy.ndimage import uniform_filter1d
    fs = SR / HOP
    sos = butter(2, [4.0, 8.0], btype="band", fs=fs, output="sos")
    B = sosfiltfilt(sos, S, axis=1)
    fm = np.zeros_like(S)
    fm[1:-1] = -(B[2:] * B[:-2])                                       # anti-phase between neighbours
    w = max(3, int(0.6 * fs))
    fm = np.maximum(uniform_filter1d(fm, w, axis=1), 0)
    norm = uniform_filter1d(S ** 2, w, axis=1) + 1e-9
    v = fm / norm
    return np.clip(v / (np.percentile(v, 99) + 1e-9), 0, 1)


def melody_contour(y_mel, monophonic=False, med_sec=None, med_amt=None, octave_fix=True, vibrato=0.0, voicing=0.25):
    """Predominant-pitch contour of a (possibly polyphonic) signal.

    Harmonic-summation salience on a 1/3-semitone CQT, then a Viterbi path that
    favours smooth pitch motion, plus a voicing decision from path salience.
    Returns (times, midi-or-nan per frame, salience per frame).
    """
    n_oct = 7
    C = np.abs(librosa.cqt(y_mel, sr=SR, hop_length=HOP, fmin=librosa.midi_to_hz(MEL_FMIN_MIDI),
                           n_bins=MEL_BPO * n_oct, bins_per_octave=MEL_BPO))
    C = (C / (C.max() + 1e-9)) ** 0.5                                  # compress dynamics
    lo = (MEL_LO - MEL_FMIN_MIDI) * 3
    hi = (MEL_HI - MEL_FMIN_MIDI) * 3
    cand = np.arange(lo, hi + 1)
    S = np.zeros((len(cand), C.shape[1]))
    for h in range(1, 9):
        off = int(round(MEL_BPO * np.log2(h)))
        idx = cand + off
        ok = idx < C.shape[0]
        S[ok] += (0.8 ** (h - 1)) * C[idx[ok]]
    if not monophonic:
        # suppress sustained accompaniment: subtract each bin's local median over ~1.5 s
        from scipy.ndimage import median_filter
        w = int((med_sec or MEL_MEDIAN_SEC) * SR / HOP) | 1
        S = np.maximum(S - (med_amt if med_amt is not None else MEL_MEDIAN_AMT) * median_filter(S, size=(1, w)), 0)
    if vibrato > 0:
        S = S * (1 + vibrato * vibrato_score(S))
    peak = S.max(axis=0)
    # Viterbi over pitch candidates: emission ∝ salience, transitions favour small steps
    P = S / (S.sum(axis=0, keepdims=True) + 1e-9)
    P = P ** 2
    P /= P.sum(axis=0, keepdims=True) + 1e-9
    trans = librosa.sequence.transition_local(len(cand), width=13, window="triangle", wrap=False)
    trans = 0.98 * trans + 0.02 / len(cand)                              # allow leaps
    trans /= trans.sum(axis=1, keepdims=True)
    path = librosa.sequence.viterbi(P, trans).astype(int)
    if octave_fix:
        # sub-octave check: if the octave below is nearly as salient, the true f0 is probably there
        below = path - MEL_BPO
        okb = below >= 0
        cols = np.arange(S.shape[1])
        bl = np.clip(below, 0, None)
        # ...but only if a pitch down there would also show its own fundamental or 3rd harmonic
        # (at f/2 and 3f/2). A high voice's 2nd, 4th, 6th harmonics alone make the octave below look
        # salient, so without this check high voices get pulled down an octave.
        h3 = int(round(MEL_BPO * np.log2(3)))
        top = np.minimum(lo + bl + h3, C.shape[0] - 1)
        evidence = np.maximum(C[lo + bl, cols], C[top, cols])
        swap = okb & (S[bl, cols] > 0.7 * S[path, cols]) & (evidence > 0.4 * C[lo + path, cols])
        path = np.where(swap, below, path)
    sal = S[path, np.arange(S.shape[1])]
    thr = voicing * np.percentile(peak, 95)
    voiced = (sal > thr) & (peak > 0.6 * voicing * np.percentile(peak, 99))
    midi = np.where(voiced, MEL_FMIN_MIDI + cand[path] / 3.0, np.nan)
    t = librosa.times_like(midi, sr=SR, hop_length=HOP)
    return t, midi, sal


def contour_onsets(t, midi, sal):
    """Note onsets = voicing starts, pitch changes, or sharp rises in melody salience."""
    db = 20 * np.log10(sal + 1e-6)
    rise = np.r_[0, np.diff(db)]
    from scipy.signal import find_peaks
    pk, _ = find_peaks(rise, height=3.0, distance=4)
    return t[pk]


def extract_melody(y_mel, beat_times, origin, subdiv, onset_times=None, monophonic=False):
    t, midi, sal = melody_contour(y_mel, monophonic)
    if onset_times is None:
        onset_times = contour_onsets(t, midi, sal)

    # build the sub-beat grid
    grid = []
    for b0, b1 in zip(beat_times[:-1], beat_times[1:]):
        for k in range(subdiv):
            grid.append((b0 + (b1 - b0) * k / subdiv, b0 + (b1 - b0) * (k + 1) / subdiv))
    cell_pitch = []
    for g0, g1 in grid:
        m = (t >= g0) & (t < g1)
        vals = midi[m]
        if m.sum() == 0 or np.isfinite(vals).mean() < 0.5:
            cell_pitch.append(None)
        else:
            cell_pitch.append(int(round(np.nanmedian(vals))))

    tol = 0.35 * (beat_times[1] - beat_times[0]) / subdiv
    onset_cell = np.zeros(len(grid), bool)
    starts = np.array([g[0] for g in grid])
    for ot in onset_times:
        i = int(np.argmin(np.abs(starts - ot)))
        if abs(starts[i] - ot) < tol:
            onset_cell[i] = True

    notes = []
    for i, p in enumerate(cell_pitch):
        sb = i / subdiv - origin
        if p is None:
            continue
        if notes and notes[-1].midi == p and cell_pitch[i - 1] == p and not onset_cell[i]:
            notes[-1].beats += 1 / subdiv
        else:
            notes.append(Note(p, sb, 1 / subdiv))
    return notes


def clean_melody(notes, key, level, subdiv):
    if not notes:
        return notes
    # remove isolated blips (single grid cell with rests around it)
    unit = 1 / subdiv
    keep = []
    for i, n in enumerate(notes):
        prev_gap = n.start_beat - (notes[i - 1].start_beat + notes[i - 1].beats) if i else 9
        next_gap = notes[i + 1].start_beat - (n.start_beat + n.beats) if i + 1 < len(notes) else 9
        if n.beats <= unit and prev_gap > 1e-6 and next_gap > 1e-6:
            continue
        keep.append(n)
    notes = keep
    # fix octave jumps relative to the previous note (typical pitch-tracker error)
    for i in range(1, len(notes)):
        d = notes[i].midi - notes[i - 1].midi
        if abs(d) > 9 and abs(d - 12 * np.sign(d)) < abs(d):
            notes[i].midi -= 12 * int(np.sign(d))
    # move the whole line into a comfortable right-hand range (median ≈ A4/B4)
    if notes:
        med = np.median([n.midi for n in notes])
        shift = 12 * int(round((70 - med) / 12))
        for n in notes:
            n.midi += shift
            while n.midi > 86:
                n.midi -= 12
            while n.midi < 57:
                n.midi += 12
    if level == "easy":
        tonic, mode = key
        scale = {(tonic + s) % 12 for s in (MAJOR_SCALE if mode == "major" else MINOR_SCALE)}
        for n in notes:
            if n.midi % 12 not in scale:
                n.midi += -1 if (n.midi - 1) % 12 in scale else 1
        # merge repeated adjacent pitches for an easier read
        merged = []
        for n in notes:
            if merged and merged[-1].midi == n.midi and abs(merged[-1].start_beat + merged[-1].beats - n.start_beat) < 1e-6:
                merged[-1].beats += n.beats
            else:
                merged.append(n)
        notes = merged
    return notes


# ──────────────────────────────────────────────────────────────── 6. analysis driver
def analyze(path, bpm=None, beats_per_bar=4, level="standard", sevenths=False,
            separate=False, verbose=True):
    log = (lambda *a: print(*a, file=sys.stderr)) if verbose else (lambda *a: None)
    log("• loading audio")
    y, y_harm, y_perc, stems = load_audio(path, separate)
    harm_src = stems.get("other", y_harm) if stems else y_harm
    if stems and "bass" in stems:
        harm_src = harm_src + stems["bass"]
    mel_src = stems.get("vocals", y_harm) if stems else y_harm

    log("• tracking beats")
    beat_times, tempo, _ = track_beats(y, y_perc, bpm)

    log("• computing chroma")
    chroma = librosa.feature.chroma_cqt(y=harm_src, sr=SR, hop_length=HOP, bins_per_octave=36)
    bass_cqt = np.abs(librosa.cqt(harm_src, sr=SR, hop_length=HOP,
                                  fmin=librosa.note_to_hz("E1"), n_bins=24))
    bass_chroma = np.zeros((12, bass_cqt.shape[1]))
    for k in range(24):
        bass_chroma[(28 + k) % 12] += bass_cqt[k]                 # E1 = midi 28
    rms = librosa.feature.rms(y=harm_src, hop_length=HOP)
    low_onset = librosa.onset.onset_strength(y=y, sr=SR, hop_length=HOP, fmax=200, n_mels=32)

    def beat_features(bts):
        cb = per_beat(chroma, bts, np.median)
        bb = per_beat(bass_chroma, bts)
        eb = per_beat(rms, bts)[0]
        lb = np.array([low_onset[min(f, len(low_onset) - 1)]
                       for f in librosa.time_to_frames(bts[:-1], sr=SR, hop_length=HOP)])
        return cb, bb, eb, lb

    chroma_b, bass_b, energy_b, lo_b = beat_features(beat_times)

    log("• estimating key")
    weights = energy_b / (energy_b.max() + 1e-9)
    key = estimate_key((chroma_b * weights).sum(axis=1) + 0.5 * bass_b.sum(axis=1) / (bass_b.sum() + 1e-9))

    log("• recognising chords + bar lines")
    labels = detect_chords(chroma_b, bass_b, energy_b, key, sevenths)
    downbeats = track_downbeats(labels, chroma_b, lo_b, beats_per_bar)
    beat_times, first_db = regrid(beat_times, downbeats, beats_per_bar)
    chroma_b, bass_b, energy_b, lo_b = beat_features(beat_times)       # second pass on bar grid
    labels = detect_chords(chroma_b, bass_b, energy_b, key, sevenths)
    first_sound = next((i for i, l in enumerate(labels) if l[0] is not None), 0)
    origin = first_db + beats_per_bar * int(np.floor((first_sound - first_db) / beats_per_bar))
    origin = max(origin, first_db % beats_per_bar)
    res = beats_per_bar if level == "easy" else (beats_per_bar // 2 if beats_per_bar % 2 == 0 else beats_per_bar)
    chords = segment_chords(labels, origin, beats_per_bar, res)

    log("• extracting melody")
    subdiv = 1 if level == "easy" else 2
    vocal_stem = bool(stems) and "vocals" in stems
    onsets = (librosa.onset.onset_detect(y=mel_src, sr=SR, hop_length=HOP, units="time")
              if vocal_stem else None)
    melody = extract_melody(mel_src, beat_times, origin, subdiv, onsets, monophonic=vocal_stem)
    melody = clean_melody(melody, key, level, subdiv)
    # trim melody to the chorded region
    if chords:
        lo_b_, hi_b_ = chords[0].start_beat, chords[-1].start_beat + chords[-1].beats
        melody = [n for n in melody if lo_b_ <= n.start_beat < hi_b_]

    title = clean_title(path)
    return Analysis(title=title, duration_sec=len(y) / SR, bpm=round(tempo, 1),
                    beats_per_bar=beats_per_bar, key_tonic=key[0], key_mode=key[1],
                    chords=chords, melody=melody, beat_times=[float(b) for b in beat_times],
                    origin_beat=origin)


def transpose_by(an: Analysis, semitones: int):
    """Shift everything by a number of semitones; the key signature follows."""
    if not semitones:
        return an
    an.key_tonic = (an.key_tonic + semitones) % 12
    an.transpose += semitones
    for c in an.chords:
        if c.root is not None:
            c.root = (c.root + semitones) % 12
    for n in an.melody:
        n.midi += semitones
    if an.melody:                                            # stay on the treble staff
        med = np.median([n.midi for n in an.melody])
        oc = -12 if med > 77 else 12 if med < 63 else 0
        for n in an.melody:
            n.midi += oc
    return an


def clean_title(path):
    t = os.path.splitext(os.path.basename(path))[0]
    t = re.sub(r"^[0-9a-f]{8}-", "", t)                                 # upload/hash prefixes
    t = re.sub(r"[_]+", " ", t)
    t = re.sub(r"[\(\[]?\b(official\s+)?(music\s+)?(video|audio|lyrics?|visualizer|hd|4k)\b[\)\]]?",
               "", t, flags=re.I)
    t = re.sub(r"\bofficial\b", "", t, flags=re.I)
    return re.sub(r"\s{2,}", " ", t).strip(" -") or "Untitled"


def transpose(an: Analysis, to_key: str):
    """Transpose analysis to a target tonic (e.g. 'C', 'Am', 'G')."""
    name = to_key.rstrip("m")
    tgt = (SHARP.index(name) if name in SHARP else FLAT.index(name))
    shift = (tgt - an.key_tonic) % 12
    if shift > 6:
        shift -= 12                                          # move the shorter way
    an.key_tonic = tgt
    an.transpose = shift
    for c in an.chords:
        if c.root is not None:
            c.root = (c.root + shift) % 12
    for n in an.melody:
        n.midi += shift
    if an.melody:                                            # keep the line in a comfy RH range
        oc = 12 * int(round((70 - np.median([n.midi for n in an.melody])) / 12))
        for n in an.melody:
            n.midi += oc
    return an


# ──────────────────────────────────────────────────────────────── 7. arrange + engrave
def split_at_bars(chords, beats_per_bar):
    """Re-strike long chords on every downbeat instead of tying across barlines."""
    out = []
    for c in chords:
        s, remaining = c.start_beat, c.beats
        while remaining > 1e-6:
            to_bar = beats_per_bar - (s % beats_per_bar)
            take = min(remaining, to_bar)
            out.append(Chord(c.root, c.quality, s, take))
            s, remaining = s + take, remaining - take
    return out


def lh_voicings(chords, level, beats_per_bar=4):
    """Return list of (start_beat, beats, [midis]) for the left hand."""
    out, prev = [], None
    for c in split_at_bars(chords, beats_per_bar):
        if c.root is None:
            continue
        ivs = QUALITIES[c.quality][0]
        root = 36 + c.root if 36 + c.root >= 41 else 48 + c.root      # F2..E3 bass root
        if level == "easy":
            notes = [root, root + 7]
        else:
            pcs = [(c.root + iv) % 12 for iv in ivs[:3]]
            cands = []
            for inv in range(3):                                      # 3 inversions near C3–G4
                order = pcs[inv:] + pcs[:inv]
                for base in (48, 52, 55):
                    v, cur = [], base - 1
                    for pc in order:
                        n = cur + 1 + ((pc - (cur + 1)) % 12)
                        v.append(n)
                        cur = n
                    if 47 <= v[0] and v[-1] <= 64:
                        cands.append(v)
            target = prev if prev else [52, 55, 60]
            notes = min(cands, key=lambda v: sum(abs(a - b) for a, b in zip(v, target)) + abs(v[0] - 52) * 0.2)
            prev = notes
        out.append((c.start_beat, c.beats, notes))
    return out


def spell(midi, flats):
    from music21 import pitch
    name = (FLAT if flats else SHARP)[midi % 12]
    return pitch.Pitch(name.replace("b", "-") + str(midi // 12 - 1))


def build_score(an: Analysis, level="standard"):
    from music21 import (chord, clef, harmony, instrument, key, layout, metadata, meter,
                         note, stream, tempo)
    flats = an.uses_flats
    bpb = an.beats_per_bar
    start = min([c.start_beat for c in an.chords] + [n.start_beat for n in an.melody] + [0.0])
    pickup_bars = int(np.ceil(-start / bpb)) if start < 0 else 0
    off0 = pickup_bars * bpb                                      # shift so all offsets ≥ 0
    end = max([c.start_beat + c.beats for c in an.chords] + [n.start_beat + n.beats for n in an.melody] + [bpb])
    total = off0 + bpb * int(np.ceil(end / bpb))

    sc = stream.Score()
    sc.insert(0, metadata.Metadata())
    sc.metadata.title = an.title
    sc.metadata.composer = f"Piano arrangement · {an.key_name} · ♩≈{int(round(an.bpm))}"

    def mkpart(cl, name):
        p = stream.Part()
        p.partName = ""
        p.partAbbreviation = ""
        p.insert(0, instrument.Piano())
        p.insert(0, cl)
        p.insert(0, key.Key(SHARP[an.key_tonic] if not flats else FLAT[an.key_tonic].replace("b", "-"),
                            an.key_mode))
        p.insert(0, meter.TimeSignature(f"{bpb}/4"))
        return p

    rh, lh = mkpart(clef.TrebleClef(), "RH"), mkpart(clef.BassClef(), "LH")
    rh.insert(0, tempo.MetronomeMark(number=int(round(an.bpm))))

    for c in an.chords:
        if c.root is None:
            continue
        fig = c.name(flats).replace("b", "-") if len(c.name(flats)) > 1 and c.name(flats)[1] == "b" else c.name(flats)
        cs = harmony.ChordSymbol(fig)
        cs.writeAsChord = False
        rh.insert(off0 + c.start_beat, cs)

    def place(part, off, dur, obj):
        part.insert(off, obj)

    for n in an.melody:
        nn = note.Note(spell(n.midi, flats), quarterLength=n.beats)
        place(rh, off0 + n.start_beat, n.beats, nn)
    for sb, beats, mids in lh_voicings(an.chords, level, bpb):
        ch = chord.Chord([spell(m, flats) for m in mids], quarterLength=beats)
        place(lh, off0 + sb, beats, ch)

    for p in (rh, lh):
        p.makeRests(fillGaps=True, inPlace=True, refStreamOrTimeRange=[0.0, float(total)])
    piano = stream.Score()
    piano.insert(0, sc.metadata)
    piano.insert(0, rh)
    piano.insert(0, lh)
    piano.insert(0, layout.StaffGroup([rh, lh], name="Piano", abbreviation="Pno.", symbol="brace",
                                      barTogether=True))
    piano = piano.makeNotation()
    return piano


def chord_chart(an: Analysis):
    flats = an.uses_flats
    bpb = an.beats_per_bar
    bars = {}
    for c in an.chords:
        b = int(np.floor(c.start_beat / bpb))
        pos = c.start_beat - b * bpb
        span = c.beats
        while span > 0:
            bars.setdefault(b, []).append((pos, c.name(flats)))
            take = bpb - pos
            span -= take
            b, pos = b + 1, 0.0
    lines, row = [], []
    for b in sorted(bars):
        cells = bars[b]
        txt = " ".join(n for i, (p, n) in enumerate(cells) if i == 0 or n != cells[i - 1][1])
        row.append(f"{txt:<10}")
        if len(row) == 4:
            lines.append("| " + " | ".join(row) + " |")
            row = []
    if row:
        lines.append("| " + " | ".join(row) + " |")
    head = (f"{an.title}\nKey: {an.key_name}   Tempo: ~{int(round(an.bpm))} BPM   "
            f"Time: {bpb}/4" + (f"   (transposed {an.transpose:+d} semitones)" if an.transpose else ""))
    return head + "\n\n" + "\n".join(lines) + "\n"


def tidy_musicxml(path):
    """music21 sometimes writes <root-alter>0</root-alter>, which renders as a natural sign."""
    with open(path) as f:
        x = f.read()
    x = re.sub(r"\s*<root-alter>0</root-alter>", "", x)
    with open(path, "w") as f:
        f.write(x)


# SMuFL glyphs Verovio emits as text; cairo has no Leipzig font, so map them to Unicode
SMUFL_TEXT = {"\uea64": "\u266d", "\uea66": "\u266f", "\uea65": "\u266e", "\ueca5": "\u2669",
              "\ued60": "\u266d", "\ued62": "\u266f", "\ued61": "\u266e"}


def _fix_svg_text(svg):
    def sub(m):
        txt = "".join(SMUFL_TEXT.get(ch, ch) for ch in m.group(2))
        size = m.group(1)
        if any(ch in "\u266d\u266f\u266e" for ch in txt):                # accidentals: match text size
            size = f"{float(size.rstrip('px')) * 0.55:g}px"
        return f'<tspan font-family="DejaVu Serif" font-size="{size}">{txt}'
    return re.sub(r'<tspan font-family="Leipzig" font-size="([^"]+)">([^<]*)', sub, svg)


def render_pdf(musicxml_path, pdf_path, png_preview=None):
    import cairosvg
    import verovio
    from pypdf import PdfWriter, PdfReader
    import io
    tk = verovio.toolkit()
    tk.setOptions({"pageWidth": 2100, "pageHeight": 2970, "scale": 45, "pageMarginTop": 80,
                   "pageMarginBottom": 80, "pageMarginLeft": 100, "pageMarginRight": 100,
                   "footer": "none", "breaks": "auto", "spacingSystem": 14,
                   "condenseFirstPage": False})
    if not tk.loadFile(musicxml_path):
        raise RuntimeError("Verovio could not load the MusicXML")
    writer = PdfWriter()
    svgs = []
    for p in range(1, tk.getPageCount() + 1):
        svg = _fix_svg_text(tk.renderToSVG(p))
        svgs.append(svg)
        writer.append(PdfReader(io.BytesIO(cairosvg.svg2pdf(bytestring=svg.encode()))))
        if p == 1 and png_preview:
            cairosvg.svg2png(bytestring=svg.encode(), write_to=png_preview, output_width=1240,
                             background_color="white")
    with open(pdf_path, "wb") as f:
        writer.write(f)
    return len(svgs)


def main(argv=None):
    ap = argparse.ArgumentParser(description="Song recording → simplified piano arrangement")
    ap.add_argument("audio")
    ap.add_argument("-o", "--outdir", default="pianoify_out")
    ap.add_argument("--level", choices=["easy", "standard"], default="standard")
    ap.add_argument("--to-key", help="transpose to this key, e.g. C, G, Am")
    ap.add_argument("--transpose", type=int, default=0, metavar="N",
                    help="transpose by N semitones (e.g. -2 or 3)")
    ap.add_argument("--bpm", type=float, help="tempo hint if beat tracking doubles/halves")
    ap.add_argument("--meter", type=int, default=4, help="beats per bar (default 4)")
    ap.add_argument("--sevenths", action="store_true", help="allow 7th chords")
    ap.add_argument("--separate", action="store_true", help="use Demucs stems if installed")
    ap.add_argument("--no-pdf", action="store_true")
    a = ap.parse_args(argv)

    an = analyze(a.audio, a.bpm, a.meter, a.level, a.sevenths, a.separate)
    if a.to_key:
        transpose(an, a.to_key)
        if a.to_key.endswith("m"):
            an.key_mode = "minor"
    transpose_by(an, a.transpose)

    os.makedirs(a.outdir, exist_ok=True)
    base = os.path.join(a.outdir, re.sub(r"[^\w\-]+", "_", an.title)[:60] or "song")
    print("• engraving score", file=sys.stderr)
    score = build_score(an, a.level)
    score.write("musicxml", fp=base + ".musicxml")
    tidy_musicxml(base + ".musicxml")
    score.write("midi", fp=base + ".mid")
    with open(base + ".chords.txt", "w") as f:
        f.write(chord_chart(an))
    d = asdict(an)
    d["key"] = an.key_name
    d["chords"] = [{**asdict(c), "name": c.name(an.uses_flats)} for c in an.chords]
    with open(base + ".analysis.json", "w") as f:
        json.dump(d, f, indent=1)
    if not a.no_pdf:
        pages = render_pdf(base + ".musicxml", base + ".pdf", base + ".preview.png")
        print(f"• {pages} page(s) of sheet music", file=sys.stderr)
    print(chord_chart(an))
    print("Outputs:", ", ".join(sorted(f for f in os.listdir(a.outdir))))
    return an


if __name__ == "__main__":
    main()