#!/usr/bin/env python3
"""Improve a MuScriptor transcription without touching the weights.

The model's own output is the starting point: it gets the notes broadly right,
then loses on three things that are visible in any difficult passage.

  1. **Timing drift.** Onsets land on the analysis hop grid, so fast passages
     smear by tens of milliseconds.
  2. **Ghost notes.** In dense or noisy passages the decoder emits notes at
     moments where nothing was struck.
  3. **Dropped notes.** Complex, blurred or quiet passages come back with gaps —
     the onset is in the audio, the note is missing from the MIDI.

All three are measurable in the audio itself, at the onsets the model already
committed to, so they can be repaired without a GPU and without retraining:
onsets are snapped to the spectral-flux peak, ghosts are dropped where the
onset has no energy, near-duplicates are merged, and gaps are filled only where
a confident pitch estimate exists.

Usage:
    python3 transcript_fix.py song.wav notes.mid out.mid
    python3 transcript_fix.py --selftest
"""

from __future__ import annotations

__version__ = "1.1"

import math
import os
import sys

import numpy as np

from muscriptor_dynamics import load_audio, measure, read_midi, write_midi

FRAME_S = 0.010
SNAP_WINDOW_S = 0.070        # how far an onset may be moved to meet a real attack
GHOST_MARGIN_DB = 9.0        # a note this far under the piece's floor is a ghost
DUPLICATE_WINDOW_S = 0.045   # same pitch this close is one note, not two
FILL_WINDOW_S = 0.085        # an onset with no note this close is a gap
FILL_MIN_CONFIDENCE = 0.45   # autocorrelation peak ratio needed to trust a pitch


# --------------------------------------------------------------------------- #
# onset detection
# --------------------------------------------------------------------------- #
def onset_envelope(samples, rate):
    """Spectral flux: positive frame-to-frame change in magnitude.

    A struck note raises every harmonic at once, so the frames where the
    magnitude jumps are the frames where something began — which is exactly what
    the model's note onsets should coincide with.
    """
    hop = max(1, int(FRAME_S * rate))
    win = 4 * hop
    if len(samples) < win + hop:
        return np.zeros(1), hop / rate
    frames = np.lib.stride_tricks.sliding_window_view(samples, win)[::hop]
    window = np.hanning(win).astype(np.float32)
    magnitude = np.abs(np.fft.rfft(frames * window, axis=1))
    flux = np.diff(magnitude, axis=0).clip(min=0).sum(axis=1)
    flux = np.concatenate([[0.0], flux])
    if flux.max() > 0:
        flux = flux / flux.max()
    return flux, hop / rate


def pick_onsets(flux, frame_s, min_gap=0.040):
    """Local maxima of the flux, with an adaptive floor."""
    if len(flux) < 3:
        return []
    floor = float(np.percentile(flux, 82))
    peaks = []
    for i in range(1, len(flux) - 1):
        if flux[i] < floor or flux[i] < flux[i - 1] or flux[i] < flux[i + 1]:
            continue
        time = i * frame_s
        if peaks and time - peaks[-1][0] < min_gap:
            if flux[i] > peaks[-1][1]:
                peaks[-1] = (time, float(flux[i]))
            continue
        peaks.append((time, float(flux[i])))
    return peaks


def refine_onset(samples, rate, time, radius=0.06):
    """Where the attack really starts, to about a millisecond.

    The flux peak is honest but late-frame: the magnitude jumps when the attack
    enters the analysis window, so its position lags the attack by roughly the
    window length. The energy derivative inside a small neighbourhood has no such
    lag, so the coarse peak is snapped to the first point where the local energy
    lifts off.
    """
    first = int(max(0, (time - radius) * rate))
    last = int(min(len(samples), (time + radius) * rate))
    if last - first < 32:
        return time
    chunk = samples[first:last]
    hop = max(1, int(0.001 * rate))
    win = max(8, int(0.006 * rate))
    if chunk.size < win + hop:
        return time
    frames = np.lib.stride_tricks.sliding_window_view(chunk, win)[::hop]
    energy = (frames * frames).mean(axis=1)
    peak = float(energy.max())
    if peak <= 0:
        return time
    # The attack is the sharpest *rise*, not the first loud frame: a window that
    # opens inside the previous note's decay starts loud and would drag the note
    # backwards.
    rise = np.diff(energy)
    if not rise.size:
        return time
    return first / rate + (int(np.argmax(rise)) + 1) * hop / rate

def estimate_pitch(chunk, rate):
    """Fundamental frequency by autocorrelation, with a confidence in 0..1.

    Used only to fill gaps the model left. A wrong pitch is worse than a missing
    note, so the caller keeps additions only when the periodic peak is strong.
    """
    if chunk.size < int(0.02 * rate):
        return None, 0.0
    signal = chunk - float(chunk.mean())
    energy = float(np.dot(signal, signal))
    if energy <= 1e-7:
        return None, 0.0
    corr = np.correlate(signal, signal, "full")[signal.size - 1:]
    low = max(2, int(rate / 2000.0))
    high = min(corr.size - 1, int(rate / 55.0))
    if high <= low + 1:
        return None, 0.0
    segment = corr[low:high]
    strongest = float(segment.max())
    if strongest <= 0 or corr[0] <= 0:
        return None, 0.0
    # The tallest peak is often a multiple of the period, which reports the note
    # an octave or two down. The first lag that reaches most of that height is
    # the period itself.
    lag = low + int(np.argmax(segment >= 0.85 * strongest))
    return rate / lag, max(0.0, min(1.0, strongest / corr[0]))


# --------------------------------------------------------------------------- #
# the repair
# --------------------------------------------------------------------------- #
def improve(notes, samples, rate, on_progress=None, on_stage=None):
    """Repair a transcription. Returns (notes, report).

    A note carries (start, dur, pitch, channel, velocity). Every stage below is
    one pass over the notes or over the onset list, so nothing here loops
    unboundedly.
    """
    def stage(percent, message):
        if on_stage:
            on_stage(percent, message)

    report = {"before": len(notes)}
    if not notes:
        return [], report

    stage(10, "детектирую онсеты в аудио")
    flux, frame_s = onset_envelope(samples, rate)
    onsets = [(refine_onset(samples, rate, time), strength)
              for time, strength in pick_onsets(flux, frame_s)]
    report["onsets"] = len(onsets)

    stage(25, "измеряю уровни нот")
    velocities = measure([note[:4] for note in notes], samples, rate, on_progress)
    levels = np.zeros(len(notes))
    for index, note in enumerate(notes):
        start = note[0]
        first = int(max(0, round(start * rate)))
        span = max(1, int(0.06 * rate))
        chunk = samples[first:first + span]
        levels[index] = 20 * math.log10(float(np.sqrt(np.mean(chunk * chunk)) + 1e-9)) if chunk.size else -120.0

    stage(45, "уточняю онсеты по атакам")
    # Each note is refined in place rather than moved to a nearby detected peak:
    # in a fast passage the neighbour's attack is also within the window, and
    # jumping to it swaps one error for a worse one.
    snapped = 0
    for index, note in enumerate(notes):
        start, dur, pitch, channel = note[:4]
        refined = refine_onset(samples, rate, start, radius=0.05)
        if abs(refined - start) > 1e-4:
            snapped += 1
        notes[index] = (refined, dur, pitch, channel, velocities[index])
    report["snapped"] = snapped

    stage(60, "убираю фантомные ноты")
    floor = float(np.percentile(levels, 12))
    kept = []
    dropped = 0
    for index, note in enumerate(notes):
        if levels[index] < floor - GHOST_MARGIN_DB:
            dropped += 1
            continue
        kept.append(note)
    notes = kept
    report["ghosts_removed"] = dropped

    stage(72, "склеиваю дубликаты")
    notes.sort(key=lambda n: (n[0], n[2]))
    merged = []
    removed = 0
    for note in notes:
        if merged and note[2] == merged[-1][2] and abs(note[0] - merged[-1][0]) < DUPLICATE_WINDOW_S:
            previous = merged[-1]
            if note[1] > previous[1]:
                merged[-1] = (previous[0], note[1], previous[2], previous[3], max(previous[4], note[4]))
            removed += 1
            continue
        merged.append(note)
    notes = merged
    report["duplicates_merged"] = removed

    stage(85, "добираю пропущенные ноты")
    # Filling is the riskiest stage: a wrong pitch is worse than a missing note.
    # Three guards keep it honest — a real onset, a confident autocorrelation, and
    # a pitch that this piece actually uses, since a performer repeats their own
    # material and noise does not. The count is capped so a noisy file cannot
    # produce more notes than it had.
    piece_pitches = sorted({note[2] for note in notes})
    budget = max(1, len(notes) // 3)
    added = 0
    window = int(0.09 * rate)
    for time, strength in onsets:
        if added >= budget:
            break
        if strength < 0.30:
            continue
        if any(abs(note[0] - time) < FILL_WINDOW_S for note in notes):
            continue
        first = int(max(0, round(time * rate)))
        chunk = samples[first:first + window]
        if chunk.size < 32:
            continue
        freq, confidence = estimate_pitch(chunk, rate)
        if freq is None or confidence < 0.60:
            continue
        pitch = int(round(69 + 12 * math.log2(freq / 440.0)))
        if piece_pitches:
            nearest = min(piece_pitches, key=lambda candidate: abs(candidate - pitch))
            if abs(nearest - pitch) > 2:
                continue
            pitch = nearest
        if not 21 <= pitch <= 108:
            continue
        level = 20 * math.log10(float(np.sqrt(np.mean(chunk * chunk)) + 1e-9))
        share = min(1.0, max(0.0, (level - floor) / max(6.0, 24.0)))
        notes.append((time, 0.18, pitch, 0, int(round(20 + share * 100))))
        added += 1
    notes.sort(key=lambda n: n[0])
    # A filled note can land on top of one that was just snapped onto the same
    # attack, which is one note heard twice. Sweep again before reporting.
    deduped = []
    for note in notes:
        if deduped and note[2] == deduped[-1][2] and abs(note[0] - deduped[-1][0]) < DUPLICATE_WINDOW_S:
            if note[1] > deduped[-1][1]:
                deduped[-1] = (deduped[-1][0], note[1], note[2], note[3], max(deduped[-1][4], note[4]))
            continue
        deduped.append(note)
    report["filled"] = added - (len(notes) - len(deduped))
    notes = deduped
    report["after"] = len(notes)

    stage(95, "готово")
    return notes, report


# --------------------------------------------------------------------------- #
# scoring
# --------------------------------------------------------------------------- #
def score(reference, candidate, tolerance=0.05):
    """Precision, recall, F1 and mean timing error, matching one-to-one."""
    used = [False] * len(candidate)
    matched = 0
    errors = []
    for start, _dur, pitch, *_rest in reference:
        best = None
        for index, other in enumerate(candidate):
            if used[index] or other[2] != pitch:
                continue
            delta = abs(other[0] - start)
            if delta <= tolerance and (best is None or delta < best[0]):
                best = (delta, index)
        if best is not None:
            used[best[1]] = True
            matched += 1
            errors.append(best[0])
    precision = matched / max(1, len(candidate))
    recall = matched / max(1, len(reference))
    f1 = 2 * precision * recall / max(1e-9, precision + recall)
    return {
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "mean_error_ms": 1000.0 * float(np.mean(errors)) if errors else float("nan"),
        "matched": matched,
    }


def _synth(rate, notes):
    total = int(rate * (max(n[0] + n[1] for n in notes) + 0.3))
    out = np.zeros(total, dtype=np.float32)
    for start, dur, pitch, amplitude in notes:
        freq = 440.0 * 2 ** ((pitch - 69) / 12)
        length = int(rate * dur)
        t = np.arange(length) / rate
        tone = sum((1.0 / k ** 1.2) * np.sin(2 * math.pi * freq * k * t) for k in (1, 2, 3, 4, 5))
        tone *= np.exp(-4.0 * t) * amplitude
        first = int(start * rate)
        out[first:first + length] += tone.astype(np.float32)
    peak = float(np.max(np.abs(out))) or 1.0
    return np.clip(out / peak * 0.95, -1, 1)


def selftest():
    rate = 22050
    # ground truth: a fast passage with a wide dynamic range, including a quiet
    # note and two notes only 90 ms apart, which is where transcription struggles
    truth = [
        (0.00, 0.22, 60, 0.55), (0.32, 0.22, 64, 0.50), (0.64, 0.22, 67, 0.45),
        (0.96, 0.22, 62, 0.24), (1.28, 0.22, 65, 0.52), (1.60, 0.22, 69, 0.48),
        (1.92, 0.22, 71, 0.14), (2.24, 0.22, 72, 0.50), (2.56, 0.22, 74, 0.46),
        (2.88, 0.22, 67, 0.22), (3.20, 0.22, 69, 0.50), (3.52, 0.22, 72, 0.44),
    ]
    samples = _synth(rate, truth)
    reference = [(start, dur, pitch, 0, 90) for start, dur, pitch, _amp in truth]

    # what a model does to this: onsets on a 20 ms grid, two notes lost in the
    # fast run, one ghost in a gap, and one note doubled
    rng = np.random.default_rng(7)
    degraded = []
    for index, (start, dur, pitch, _channel, _vel) in enumerate(reference):
        if index in (3, 6, 9):                   # lost notes, three quiet ones among them
            continue
        jitter = 0.02 * (index % 3 - 1) + rng.normal(0, 0.004)
        degraded.append((max(0.0, start + jitter), dur, pitch, 0, 64))
    degraded.append((3.95, 0.2, 66, 0, 64))     # ghost, in real silence
    degraded.append((0.11, 0.22, 64, 0, 64))    # duplicate of the second note
    degraded.sort()

    for probe in (0.96, 1.92):
        first = int(probe * rate)
        freq, confidence = estimate_pitch(samples[first:first + int(0.09 * rate)], rate)
        print("  зонд %.2f с: f0=%.0f Гц (pitch %d), conf %.2f"
              % (probe, freq or 0, round(69 + 12 * math.log2((freq or 440) / 440)), confidence), flush=True)
    before = score(reference, degraded)
    print("до правки:  P=%.2f R=%.2f F1=%.2f, ошибка онсета %.0f мс, нот %d"
          % (before["precision"], before["recall"], before["f1"],
             before["mean_error_ms"], len(degraded)), flush=True)

    improved, report = improve(degraded, samples, rate,
                               on_stage=lambda p, m: print("  [%3d%%] %s" % (p, m), flush=True))
    after = score(reference, improved)
    print("после:      P=%.2f R=%.2f F1=%.2f, ошибка онсета %.0f мс, нот %d"
          % (after["precision"], after["recall"], after["f1"],
             after["mean_error_ms"], len(improved)), flush=True)
    print("отчёт: %s" % report, flush=True)

    assert after["mean_error_ms"] < before["mean_error_ms"] * 0.6, "тайминг не улучшился"
    assert report["ghosts_removed"] >= 1, "фантом не убран"
    assert report["snapped"] >= len(degraded) * 0.6, "онсеты почти не уточнились"
    assert after["f1"] >= before["f1"] - 0.05, "F1 просел больше чем на 0.05"
    print("selftest: OK", flush=True)


def run(audio_path, midi_path, out_path):
    print("[  0%] читаю аудио и MIDI", flush=True)
    samples, rate = load_audio(audio_path)
    notes, tempo, _division = read_midi(midi_path)
    bpm = 60000000.0 / tempo[0][1] if tempo else 120.0
    print("[ 10%] нот во входе: %d" % len(notes), flush=True)
    improved, report = improve(notes, samples, rate,
                               on_stage=lambda p, m: print("[%3d%%] %s" % (p, m), flush=True))
    write_midi(out_path, improved, bpm)
    print("[100%] записан %s: %s" % (out_path, report), flush=True)
    return improved


if __name__ == "__main__":
    if "--selftest" in sys.argv:
        selftest()
    elif len(sys.argv) == 4:
        run(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print(__doc__)
        raise SystemExit(2)