#!/usr/bin/env python3
"""Dynamics for MuScriptor output: loudness becomes velocity.

MuScriptor Large finds pitch, onset, duration and instrument, but it does not
predict velocity, so every note in its MIDI arrives at the same strength and a
forte passage reads exactly like a piano one. The strength is measurable in the
recording itself, note by note, at the onset the model already placed: what is
missing is not a better model but a measurement.

So this takes the audio and the model's MIDI, measures each note, and writes a
MIDI whose velocities follow the performance. Nothing is retrained and nothing
loops: one pass over the notes, progress printed in tens of percent.

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

from __future__ import annotations

__version__ = "1.1"

import math
import os
import struct
import sys
import wave

import numpy as np

# How far into a note the loudness is measured. A piano hammer peaks within a
# few tens of milliseconds; a bowed note swells for longer. A quarter of a second
# covers both without wandering into the next note in a fast passage.
ONSET_WINDOW_S = 0.25
# The dynamic range mapped onto 1..127, in decibels. Everything at or below the
# floor is a whisper, everything at or above the ceiling is a full strike.
FLOOR_DB = -34.0
CEIL_DB = 0.0
# Gamma below 1 lifts quiet notes so a soft passage stays playable rather than
# collapsing into single-digit velocities.
CURVE = 0.65


# --------------------------------------------------------------------------- #
# audio
# --------------------------------------------------------------------------- #
def load_audio(path):
    """Mono float32 samples and the sample rate. WAV natively, MP3/other via ffmpeg."""
    if path.lower().endswith(".wav"):
        with wave.open(path, "rb") as w:
            channels = w.getnchannels()
            width = w.getsampwidth()
            rate = w.getframerate()
            frames = w.readframes(w.getnframes())
    else:
        import subprocess
        import tempfile

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            out = tmp.name
        try:
            subprocess.run(
                ["ffmpeg", "-v", "error", "-y", "-i", path, "-ac", "1",
                 "-ar", "22050", "-f", "wav", out],
                check=True,
            )
            with wave.open(out, "rb") as w:
                channels = w.getnchannels()
                width = w.getsampwidth()
                rate = w.getframerate()
                frames = w.readframes(w.getnframes())
        finally:
            try:
                os.unlink(out)
            except OSError:
                pass

    if width == 2:
        data = np.frombuffer(frames, dtype="<i2").astype(np.float32) / 32768.0
    elif width == 4:
        data = np.frombuffer(frames, dtype="<i4").astype(np.float32) / 2147483648.0
    elif width == 1:
        data = (np.frombuffer(frames, dtype=np.uint8).astype(np.float32) - 128.0) / 128.0
    else:
        raise SystemExit("unsupported sample width: " + str(width))

    if channels > 1:
        data = data.reshape(-1, channels).mean(axis=1)
    return data, rate


# --------------------------------------------------------------------------- #
# MIDI reading and writing (Standard MIDI File, the parts a transcription uses)
# --------------------------------------------------------------------------- #
MARKER = b"muscriptor:bar_offset="


def _vlq(data, i):
    value = 0
    while True:
        byte = data[i]
        i += 1
        value = (value << 7) | (byte & 0x7F)
        if not byte & 0x80:
            return value, i


def _put_vlq(out, value):
    chunk = [value & 0x7F]
    value >>= 7
    while value:
        chunk.append((value & 0x7F) | 0x80)
        value >>= 7
    out.extend(reversed(chunk))


def read_midi(path):
    """Notes as (start_s, dur_s, pitch, channel) plus the raw tempo map."""
    raw = open(path, "rb").read()
    if not raw.startswith(b"MThd"):
        raise SystemExit("not a MIDI file: " + path)
    division = struct.unpack(">H", raw[12:14])[0]
    if division & 0x8000:
        raise SystemExit("SMPTE timecode MIDI is not supported")
    pos = 14
    tempo = [(0, 500000)]
    sounding = {}
    notes = []
    while pos + 8 <= len(raw) and raw[pos:pos + 4] == b"MTrk":
        length = struct.unpack(">I", raw[pos + 4:pos + 8])[0]
        body = raw[pos + 8:pos + 8 + length]
        pos += 8 + length
        tick = 0
        status = None
        i = 0
        while i < len(body):
            delta, i = _vlq(body, i)
            tick += delta
            if i >= len(body):
                break
            byte = body[i]
            if byte == 0xFF:
                i += 1
                meta = body[i]
                i += 1
                size, i = _vlq(body, i)
                if meta == 0x51 and size == 3:
                    tempo.append((tick, (body[i] << 16) | (body[i + 1] << 8) | body[i + 2]))
                i += size
                status = None
            elif byte in (0xF0, 0xF7):
                i += 1
                size, i = _vlq(body, i)
                i += size
                status = None
            elif byte & 0x80:
                status = byte
                i += 1
                high = status & 0xF0
                channel = status & 0x0F
                count = 1 if high in (0xC0, 0xD0) else 2
                args = body[i:i + count]
                i += count
                if high == 0x90 and args[1] > 0:
                    sounding.setdefault((channel, args[0]), []).append(tick)
                elif high == 0x80 or (high == 0x90 and args[1] == 0):
                    stack = sounding.get((channel, args[0]))
                    if stack:
                        start = stack.pop(0)
                        notes.append([start, tick - start, args[0], channel])
            else:
                if status is None:
                    break
                high = status & 0xF0
                channel = status & 0x0F
                count = 1 if high in (0xC0, 0xD0) else 2
                args = body[i:i + count]
                i += count
                if high == 0x90 and args[1] > 0:
                    sounding.setdefault((channel, args[0]), []).append(tick)
                elif high == 0x80 or (high == 0x90 and args[1] == 0):
                    stack = sounding.get((channel, args[0]))
                    if stack:
                        start = stack.pop(0)
                        notes.append([start, tick - start, args[0], channel])

    for (channel, pitch), stack in sounding.items():
        for start in stack:
            notes.append([start, 0, pitch, channel])

    tempo.sort(key=lambda item: item[0])

    def to_seconds(tick):
        seconds = 0.0
        last = 0
        us_per_quarter = 500000
        for at, us in tempo:
            if at > tick:
                break
            seconds += (at - last) / division * (us_per_quarter / 1000000.0)
            last = at
            us_per_quarter = us
        seconds += (tick - last) / division * (us_per_quarter / 1000000.0)
        return seconds

    converted = [(to_seconds(s), max(0.05, to_seconds(s + d) - to_seconds(s)), p, c)
                 for s, d, p, c in notes]
    converted.sort(key=lambda item: item[0])
    return converted, tempo, division


def write_midi(path, notes, bpm=120):
    """A single-track MIDI at a constant tempo, velocities as measured.

    The bar-offset marker the box writes is deliberately not reproduced: it is a
    text event that anything rendering MIDI text shows as garbage, and the
    alignment it carries is only meaningful to the box's own auralizer.
    """
    division = 480
    us_per_quarter = int(round(60000000.0 / max(1.0, bpm)))
    events = [(0, b"\xff\x51\x03" + us_per_quarter.to_bytes(3, "big"))]
    events.append((0, b"\xff\x03\x08pianoify"))
    for start, dur, pitch, channel, velocity in notes:
        on_tick = int(round(start * division * 1000000.0 / us_per_quarter))
        off_tick = int(round((start + dur) * division * 1000000.0 / us_per_quarter))
        channel = max(0, min(15, channel))
        events.append((on_tick, bytes([0x90 | channel, pitch, max(1, min(127, velocity))])))
        events.append((off_tick, bytes([0x80 | channel, pitch, 0])))
    events.sort(key=lambda item: (item[0], item[1][0] & 0xF0 == 0x80))

    track = bytearray()
    last = 0
    for tick, payload in events:
        _put_vlq(track, tick - last)
        track.extend(payload)
        last = tick
    _put_vlq(track, division * 2)
    track.extend(b"\xff\x2f\x00")

    blob = bytearray(b"MThd" + struct.pack(">IHHH", 6, 0, 1, division))
    blob += b"MTrk" + struct.pack(">I", len(track)) + track
    with open(path, "wb") as f:
        f.write(blob)
    return path


# --------------------------------------------------------------------------- #
# the measurement
# --------------------------------------------------------------------------- #
def measure(notes, samples, rate, progress=None):
    """Velocity per note, from the loudness and the attack at its onset.

    Two things make a note read as fortissimo rather than piano: how loud it is
    overall, and how hard it starts. A drummer and a pianist both hit harder for
    an accent, so the attack-to-body ratio is blended in — it is what separates a
    struck note from a sustained one at the same level.
    """
    levels = np.zeros(len(notes))
    attacks = np.zeros(len(notes))
    window = max(1, int(ONSET_WINDOW_S * rate))
    total = max(1, len(notes))
    for index, (start, dur, pitch, _channel) in enumerate(notes):
        first = int(max(0, round(start * rate)))
        span = min(window, max(1, int(max(0.03, min(dur, ONSET_WINDOW_S)) * rate)))
        chunk = samples[first:first + span]
        if chunk.size == 0:
            levels[index] = -120.0
            attacks[index] = 0.0
        else:
            rms = float(np.sqrt(np.mean(chunk * chunk)) + 1e-9)
            peak = float(np.max(np.abs(chunk)) + 1e-9)
            levels[index] = 20.0 * math.log10(rms)
            attacks[index] = peak / rms
        if progress and index % max(1, total // 10) == 0:
            progress("измеряю ноты: %d%%" % (100 * index // total))

    if not len(notes):
        return []

    loud = np.percentile(levels, 99.7)
    quiet = np.percentile(levels, 4)
    if loud - quiet < 6.0:                     # a flat piece: keep gentle dynamics
        loud = quiet + 6.0
    attack_med = float(np.median(attacks)) or 1.0
    attack_span = max(1e-6, float(np.percentile(attacks, 90) - np.percentile(attacks, 10)))

    velocities = []
    for index in range(len(notes)):
        share = (levels[index] - quiet) / (loud - quiet)
        share = min(1.0, max(0.0, share))
        accent = (attacks[index] - (attack_med - attack_span / 2)) / attack_span
        accent = min(1.0, max(-1.0, accent))
        # Multiplicative, not additive: an accent should sharpen the note's own
        # level rather than lift every loud note onto the same ceiling.
        shaped = min(1.0, max(0.0, share ** CURVE * (1.0 + 0.18 * accent)))
        # A floor of 12: velocity 1 is inaudible on every synth, so the quietest
        # measured note still has to be playable, and the loudest keeps headroom.
        velocities.append(int(round(12 + shaped * 112)))
    return velocities


def histogram(velocities):
    buckets = [0] * 6
    for velocity in velocities:
        buckets[min(5, velocity * 6 // 128)] += 1
    return buckets


def run(audio_path, midi_path, out_path):
    steps = 5

    def tell(percent, what):
        print("[%3d%%] %s" % (percent, what), flush=True)

    tell(100 * 0 // steps, "читаю аудио " + os.path.basename(audio_path))
    samples, rate = load_audio(audio_path)
    tell(100 * 1 // steps, "аудио: %.1f с, %d Гц" % (len(samples) / rate, rate))

    notes, tempo, division = read_midi(midi_path)
    if not notes:
        raise SystemExit("no notes in " + midi_path)
    tell(100 * 2 // steps, "MIDI: %d нот, первый темп %d мкс/четверть"
         % (len(notes), tempo[0][1]))

    def progress(message):
        print("        " + message, flush=True)

    velocities = measure(notes, samples, rate, progress)
    tell(100 * 3 // steps, "динамика: velocity %d..%d, медиана %d"
         % (min(velocities), max(velocities), int(np.median(velocities))))

    bpm = 60000000.0 / tempo[0][1]
    out_notes = [(start, dur, pitch, channel, velocities[i])
                 for i, (start, dur, pitch, channel) in enumerate(notes)]
    write_midi(out_path, out_notes, bpm)
    tell(100 * 4 // steps, "записан " + os.path.basename(out_path))

    buckets = histogram(velocities)
    labels = ["pp", "p", "mp", "mf", "f", "ff"]
    print("        динамические оттенки: "
          + ", ".join(labels[i] + "=" + str(buckets[i]) for i in range(6)), flush=True)
    tell(100, "готово")
    return velocities


# --------------------------------------------------------------------------- #
# self-test: the measurement has to follow the loudness that is really there
# --------------------------------------------------------------------------- #
def selftest():
    rate = 22050
    amplitudes = [0.08, 0.5, 0.12, 0.45, 0.09, 0.55, 0.15, 0.4]
    pitches = [60, 62, 64, 65, 67, 69, 71, 72]
    note_seconds = 0.4
    total = int(rate * (note_seconds * len(pitches) + 0.2))
    signal = np.zeros(total, dtype=np.float32)
    for index, (amp, pitch) in enumerate(zip(amplitudes, pitches)):
        freq = 440.0 * 2 ** ((pitch - 69) / 12)
        length = int(rate * note_seconds)
        t = np.arange(length) / rate
        tone = sum((1.0 / k) * np.sin(2 * math.pi * freq * k * t) for k in (1, 2, 3))
        tone *= np.exp(-6.0 * t)
        start = int(index * note_seconds * rate)
        signal[start:start + length] += (amp * tone).astype(np.float32)
    signal = np.clip(signal / max(1e-9, np.max(np.abs(signal))) * 0.95, -1, 1)

    here = os.path.dirname(os.path.abspath(__file__))
    wav_path = os.path.join(here, "selftest_dynamics.wav")
    midi_path = os.path.join(here, "selftest_dynamics.mid")
    out_path = os.path.join(here, "selftest_dynamics_out.mid")

    with wave.open(wav_path, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(rate)
        w.writeframes((signal * 32767).astype("<i2").tobytes())

    flat = [(index * note_seconds, note_seconds * 0.9, pitch, 0, 64)
            for index, pitch in enumerate(pitches)]
    write_midi(midi_path, flat, 120)

    print("self-test: 8 нот, амплитуды " + ", ".join("%.2f" % a for a in amplitudes), flush=True)
    velocities = run(wav_path, midi_path, out_path)

    samples, load_rate = load_audio(wav_path)
    notes, _tempo, _division = read_midi(out_path)
    assert load_rate == rate, "sample rate changed"
    assert len(notes) == len(pitches), "note count changed: %d" % len(notes)

    get = dict()
    for start, _dur, pitch, _channel in notes:
        get[round(start / note_seconds)] = pitch
    del get

    soft = [velocities[i] for i in (0, 2, 4)]
    loud = [velocities[i] for i in (1, 3, 5)]
    correlation = float(np.corrcoef(amplitudes[:len(velocities)], velocities)[0][1])
    print("        тихие: %s | громкие: %s" % (soft, loud), flush=True)
    print("        корреляция амплитуда/velocity: %.3f" % correlation, flush=True)

    assert min(loud) - max(soft) > 20, "динамика не разделилась"
    assert correlation > 0.85, "velocity не следует за громкостью"
    assert all(1 <= v <= 127 for v in velocities), "velocity вне диапазона"
    print("self-test: OK", flush=True)
    for path in (wav_path, midi_path, out_path):
        try:
            os.unlink(path)
        except OSError:
            pass


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)