"""Живая громкость вместо константы 100.

MuScriptor предсказывает только «нота началась», а velocity в MIDI пишет
одну и ту же — 100. Здесь громкость берётся из самой записи: для каждой ноты
меряется RMS и пик в окне 250 мс от атаки, раскладывается по перцентилям
и превращается в velocity 1..124.

Ключевое: файл не пересочиняется. Меняются ТОЛЬКО байты velocity у note_on,
поэтому каналы, программа, темп, метки и всё прочее остаётся как было.

Проверка: python plus/velocity.py audio.wav in.mid out.mid
"""

import math
import struct
import sys


# ---------- разбор и правка SMF ----------

def _vlq(data, index):
    value, shift = 0, 0
    while True:
        byte = data[index]
        index += 1
        value = (value << 7) | (byte & 0x7F)
        if not byte & 0x80:
            return value, index
        shift += 7
        if shift > 28:
            raise ValueError("битый delta")


def scan_midi(data):
    """Вернуть ноты и места, где лежит их velocity."""
    if data[:4] != b"MThd":
        raise ValueError("это не MIDI")
    header_len = struct.unpack(">I", data[4:8])[0]
    division = struct.unpack(">H", data[8 + header_len - 2:8 + header_len])[0]
    ticks_per_beat = division if division & 0x8000 == 0 else 480

    notes = []
    tempos = [(0, 500000)]
    index = 8 + header_len
    while index + 8 <= len(data):
        chunk = data[index:index + 4]
        length = struct.unpack(">I", data[index + 4:index + 8])[0]
        body_start, body_end = index + 8, index + 8 + length
        if chunk != b"MTrk":
            index = body_end
            continue
        i, tick, running = body_start, 0, None
        open_notes = {}
        while i < body_end:
            delta, i = _vlq(data, i)
            tick += delta
            status = data[i]
            if status < 0x80:
                if running is None:
                    raise ValueError("running status без статуса")
                status = running
            else:
                i += 1
                if status < 0xF0:
                    running = status
            kind = status & 0xF0
            if status == 0xFF:
                meta_type = data[i]
                meta_len, i = _vlq(data, i + 1)
                if meta_type == 0x51 and meta_len >= 3:
                    tempos.append((tick, int.from_bytes(data[i:i + 3], "big")))
                i += meta_len
            elif status in (0xF0, 0xF7):
                meta_len, i = _vlq(data, i)
                i += meta_len
            elif kind in (0x80, 0x90):
                pitch, velocity = data[i], data[i + 1]
                key = (status & 0x0F, pitch)
                if kind == 0x90 and velocity > 0:
                    open_notes[key] = len(notes)
                    notes.append({"tick": tick, "pitch": pitch, "channel": status & 0x0F,
                                  "velocity_at": i + 1, "velocity": velocity,
                                  "track": body_start})
                elif key in open_notes:
                    notes[open_notes.pop(key)]["end_tick"] = tick
                i += 2
            elif kind in (0xA0, 0xB0, 0xE0):
                i += 2
            elif kind in (0xC0, 0xD0):
                i += 1
            else:
                raise ValueError("неизвестный статус 0x%02X" % status)
        index = body_end

    for note in notes:
        note.setdefault("end_tick", note["tick"] + ticks_per_beat // 4)

    tempos.sort()
    notes.sort(key=lambda n: n["tick"])
    for note in notes:
        note["start"] = ticks_to_seconds(note["tick"], tempos, ticks_per_beat)
        note["end"] = ticks_to_seconds(note["end_tick"], tempos, ticks_per_beat)
    return notes


def ticks_to_seconds(tick, tempos, ticks_per_beat):
    seconds, previous_tick, micros = 0.0, 0, 500000
    for tempo_tick, tempo_micros in tempos:
        if tick <= tempo_tick:
            break
        seconds += (tempo_tick - previous_tick) / ticks_per_beat * micros / 1_000_000.0
        previous_tick, micros = tempo_tick, tempo_micros
    seconds += (tick - previous_tick) / ticks_per_beat * micros / 1_000_000.0
    return seconds


# ---------- громкость из записи ----------

def decode_audio(path):
    import numpy
    try:
        import soundfile
        samples, rate = soundfile.read(path, dtype="float32", always_2d=True)
        return samples.mean(axis=1), rate
    except Exception:
        pass
    try:
        import librosa
        samples, rate = librosa.load(path, sr=None, mono=True)
        return numpy.asarray(samples, dtype="float32"), rate
    except Exception as problem:
        raise RuntimeError("не читается аудио: %s" % problem)


def measure(notes, samples, rate, window=0.25):
    import numpy
    window_n = max(16, int(window * rate))
    levels, attacks = [], []
    for note in notes:
        start = int(note["start"] * rate)
        if start >= len(samples):
            continue
        chunk = samples[start:start + window_n]
        if len(chunk) < 16:
            continue
        rms = float(numpy.sqrt(numpy.mean(chunk.astype("float64") ** 2)))
        peak = float(numpy.max(numpy.abs(chunk)))
        if rms <= 1e-7:
            continue
        levels.append(20.0 * math.log10(rms))
        attacks.append(max(1e-6, peak / rms))
        note["_level"] = levels[-1]
        note["_attack"] = attacks[-1]
    return levels, attacks


def to_velocity(notes):
    import numpy
    levels = numpy.array([n["_level"] for n in notes])
    attacks = numpy.array([n["_attack"] for n in notes])
    quiet = numpy.percentile(levels, 4)
    loud = numpy.percentile(levels, 99.7)
    span = max(0.5, loud - quiet)
    low_attack = numpy.percentile(attacks, 10)
    high_attack = numpy.percentile(attacks, 90)
    attack_span = max(0.05, high_attack - low_attack)
    for note in notes:
        share = min(1.0, max(0.0, (note["_level"] - quiet) / span))
        accent = min(1.0, max(0.0, (note["_attack"] - low_attack) / attack_span))
        shaped = (share ** 0.65) * (1.0 + 0.18 * accent)
        note["_new"] = int(round(1 + min(1.0, shaped) * 123))
    return [n["_new"] for n in notes]


# ---------- главное ----------

def apply(audio_path, midi_bytes):
    """Вернуть MIDI с измеренной громкостью и краткую статистику."""
    notes = scan_midi(midi_bytes)
    if not notes:
        return midi_bytes, {"notes": 0}
    samples, rate = decode_audio(audio_path)
    measured = [n for n in notes if n.get("_level") is not None] or notes
    measure(notes, samples, rate)
    measured = [n for n in notes if n.get("_level") is not None]
    if not measured:
        return midi_bytes, {"notes": len(notes), "measured": 0}
    velocities = to_velocity(measured)

    fallback = int(round(sum(velocities) / len(velocities)))
    patched = bytearray(midi_bytes)
    silent = 0
    for note in notes:
        if note.get("_new") is None:
            note["_new"] = fallback          # окно попало в тишину
            silent += 1
        patched[note["velocity_at"]] = note["_new"]
    stats = {
        "notes": len(notes),
        "measured": len(measured),
        "velocity_min": min(velocities),
        "velocity_max": max(velocities),
        "velocity_mean": round(sum(velocities) / len(velocities), 1),
        "distinct": len({n["_new"] for n in notes}),
        "silent": silent,
        "duration_s": round(notes[-1]["end"], 2),
    }
    return bytes(patched), stats


def main(argv=None):
    argv = argv or sys.argv[1:]
    if len(argv) < 3:
        print(__doc__)
        return 2
    audio_path, midi_path, out_path = argv[:3]
    with open(midi_path, "rb") as handle:
        original = handle.read()
    patched, stats = apply(audio_path, original)
    with open(out_path, "wb") as handle:
        handle.write(patched)
    print("нот: %(notes)d, измерено: %(measured)d | velocity %(velocity_min)d..%(velocity_max)d"
          " (среднее %(velocity_mean).1f, разных значений %(distinct)d)" % stats)
    return 0


if __name__ == "__main__":
    sys.exit(main())