"""Проверка кандидатов по аудио: есть ли в записи энергия на этой высоте.

Модель с узким conditioning начинает выдумывать партии (на тесте 400 нот баса
на миксе, где бас есть, но не в таких высотах). Поэтому каждый кандидат
проверяется по звуку: в окне от атаки считается энергия на основной частоте
(Гёрцель) и сравнивается с медианной энергией по соседним полутонам. Нота
выживает, если её тон реально выделяется.
"""

import math
import wave

import numpy


def load_wav(path):
    with wave.open(path, "rb") as handle:
        rate = handle.getframerate()
        channels = handle.getnchannels()
        width = handle.getsampwidth()
        frames = handle.readframes(handle.getnframes())
    if width != 2:
        raise ValueError("поддерживается 16-битный PCM")
    data = numpy.frombuffer(frames, dtype="<i2").astype("float32") / 32768.0
    if channels > 1:
        data = data.reshape(-1, channels).mean(axis=1)
    return data, rate


def goertzel(samples, rate, freq):
    """Энергия на одной частоте без полного FFT."""
    n = len(samples)
    if n < 32 or freq <= 0:
        return 0.0
    k = int(0.5 + n * freq / rate)
    omega = 2.0 * math.pi * k / n
    coeff = 2.0 * math.cos(omega)
    s_prev = s_prev2 = 0.0
    for value in samples:
        s = value + coeff * s_prev - s_prev2
        s_prev2, s_prev = s_prev, s
    power = s_prev2 ** 2 + s_prev ** 2 - coeff * s_prev * s_prev2
    return max(0.0, power) / (n * n)


def pitch_energy(samples, rate, midi):
    freq = 440.0 * 2 ** ((midi - 69) / 12.0)
    return goertzel(samples, rate, freq) + goertzel(samples, rate, freq * 2)


def verify(notes, audio_path, window=0.25, neighbours=3, threshold=2.0):
    """Отметить ноты как подтверждённые записью. Возвращает (kept, dropped, stats)."""
    samples, rate = load_wav(audio_path)
    length = max(64, int(window * rate))
    kept, dropped = [], []
    for note in notes:
        if note.get("instrument") == "drums" or note.get("is_drum"):
            note["energy_ratio"] = None      # у барабанов нет основного тона
            kept.append(note)
            continue
        start = int(note["start"] * rate)
        if start >= len(samples):
            dropped.append(note)
            continue
        chunk = samples[start:start + length]
        if len(chunk) < 64:
            dropped.append(note)
            continue
        mine = pitch_energy(chunk, rate, note["pitch"])
        reference = []
        for step in range(1, neighbours + 1):
            for direction in (-1, 1):
                reference.append(pitch_energy(chunk, rate, note["pitch"] + direction * step))
        floor = float(numpy.median(reference)) + 1e-9
        ratio = mine / floor
        note["energy_ratio"] = round(ratio, 2)
        if ratio >= threshold or note.get("support", 1) >= 2:
            kept.append(note)
        else:
            dropped.append(note)
    stats = {
        "kept": len(kept),
        "dropped": len(dropped),
        "threshold": threshold,
    }
    pitched = [n["energy_ratio"] for n in kept if n.get("energy_ratio") is not None]
    if pitched:
        stats["median_ratio"] = round(float(numpy.median(pitched)), 2)
    return kept, dropped, stats