#!/usr/bin/env python3
"""Проверка акустики без слуха (#96): энергия по полосам в дорожке с преградой и без неё.

    python3 tools/acoustics_check.py open.raw wall.raw            # s16le, 48000 Гц, стерео
    python3 tools/acoustics_check.py open.raw wall.raw --rate 48000 --channels 2

Дорожки берутся из `Audio::film_mix` (экспорт видео пишет `<имя>.audio.raw`) или из теста
`cargo test -p horizon-client --lib audio::tests::export_wall_tracks -- --ignored` (HORIZON_ACOUSTICS_OUT=папка).
Для каждой дорожки считается энергия левого канала в четырёх полосах (до 300 Гц, 300-1500, 1500-5000, выше
5000) и разница в дБ между дорожками. Преграда должна срезать верхи (выше 5 кГц не меньше чем на 15 дБ) и
оставить низы (до 300 Гц теряют не больше 8 дБ). Коды выхода: 0 — в порогах, 1 — нет, 2 — ошибка аргументов.
Режим `--echo` (стадия 2): в дорожке с одним щелчком ищутся пики огибающей после прямого звука, их задержки
(мс от прямого) сравниваются с ожидаемыми от геометрии:

    python3 tools/acoustics_check.py --echo gorge.raw --expect 187.3,254.1 [--tol 1.0]
    python3 tools/acoustics_check.py --echo field.raw --expect none
    python3 tools/acoustics_check.py --echo-dir папка        # gorge/field/room и expected.txt из export_echo_tracks

Дорожки пишет `audio::tests::export_echo_tracks` (HORIZON_ACOUSTICS_OUT=папка). Коды выхода те же.
Режим `--reverb-dir` (стадия 3): папка из `audio::tests::export_reverb_tracks` (щелчок в комнате, тоннеле, ущелье
и поле; `expected_reverb.txt`: предсказанный спад и уровень из замера пространства). По хвосту после прямого звука
считаются RT60 (обратное интегрирование Шрёдера, спад с -5 до -35 дБ, удвоенный) и энергия хвоста к энергии щелчка:

    python3 tools/acoustics_check.py --reverb-dir папка [--rt-tol 0.35]

Проверяется: в поле хвоста нет (энергия хвоста меньше -90 дБ от щелчка), спад упорядочен поле < ущелье < тоннель <
комната, у комнаты и тоннеля измеренный RT60 в пределах `--rt-tol` от того, что дают их посылки в слоты (колонка
`RT60 expected`: сумма экспонент слотов, спад -5...-35 дБ; слой из двух слотов читается длиннее короткого).
Только стандартная библиотека.
"""
import argparse
import array
import cmath
import math
import sys

BANDS = [("<300", 0, 300), ("300-1500", 300, 1500), ("1500-5000", 1500, 5000), (">5000", 5000, 1e9)]


def fft(a):
    n = len(a)
    if n == 1:
        return a
    even, odd = fft(a[0::2]), fft(a[1::2])
    t = [cmath.exp(-2j * math.pi * k / n) * odd[k] for k in range(n // 2)]
    return [even[k] + t[k] for k in range(n // 2)] + [even[k] - t[k] for k in range(n // 2)]


def bands(path, rate, channels):
    raw = array.array("h")
    with open(path, "rb") as f:
        data = f.read()
    raw.frombytes(data[: len(data) // 2 * 2])
    left = [s / 32768.0 for s in raw[0::channels]]
    n = 1
    while n * 2 <= len(left):
        n *= 2
    n = min(n, 1 << 16)
    if n < 1024:
        sys.exit(f"{path}: too short ({len(left)} frames)")
    # A Hann window over the loudest n frames' start: the burst is at the start of the track.
    w = [left[i] * (0.5 - 0.5 * math.cos(2 * math.pi * i / (n - 1))) for i in range(n)]
    spec = fft([complex(v) for v in w])
    out = [0.0] * len(BANDS)
    for k in range(1, n // 2):
        hz = k * rate / n
        for i, (_, lo, hi) in enumerate(BANDS):
            if lo <= hz < hi:
                out[i] += abs(spec[k]) ** 2
    return out


def echo_peaks(path, rate, channels, min_gap_ms=10.0, rel=0.02):
    """Delays (ms from the first peak) of the peaks of the smoothed envelope of a track."""
    raw = array.array("h")
    with open(path, "rb") as f:
        data = f.read()
    raw.frombytes(data[: len(data) // 2 * 2])
    frames = [sum(abs(raw[i + c]) for c in range(channels)) / 32768.0 for i in range(0, len(raw) - channels + 1, channels)]
    win = max(1, rate // 1000)  # 1 ms
    env = []
    acc = 0.0
    for i, v in enumerate(frames):
        acc += v
        if i >= win:
            acc -= frames[i - win]
        env.append(acc / win)
    first = max(range(min(len(env), rate // 20)), key=env.__getitem__)  # the direct sound: loudest in the first 50 ms
    top = env[first]
    gap = int(min_gap_ms * rate / 1000)
    peaks, i = [], first + int(0.02 * rate)
    while i < len(env):
        if env[i] >= rel * top:
            j = max(range(i, min(len(env), i + gap)), key=env.__getitem__)
            peaks.append((j - first) * 1000.0 / rate)
            i = j + gap
        else:
            i += 1
    return peaks


def check_echo(path, expect, tol, rate, channels):
    peaks = echo_peaks(path, rate, channels)
    print(f"{path}: peaks at {', '.join(f'{p:.1f}' for p in peaks) or 'none'} ms; expected {', '.join(f'{e:.1f}' for e in expect) or 'none'} ms")
    ok = len(peaks) == len(expect) and all(abs(p - e) <= tol for p, e in zip(peaks, expect))
    print(f"  {'ok' if ok else 'FAIL'} (tolerance {tol} ms)")
    return ok


def reverb_track(path, rate, channels, dry_ms=60.0):
    """(tail energy relative to the click in dB or None, RT60 in s or None) of a click's track."""
    raw = array.array("h")
    with open(path, "rb") as f:
        data = f.read()
    raw.frombytes(data[: len(data) // 2 * 2])
    left = [s / 32768.0 for s in raw[0::channels]]
    cut = int(dry_ms * rate / 1000)
    dry = sum(v * v for v in left[:cut])
    tail = left[cut:]
    energy = sum(v * v for v in tail)
    if dry <= 0 or energy <= 0 or 10 * math.log10(energy / dry) < -90:
        return None, None
    edc, acc = [0.0] * len(tail), 0.0
    for i in range(len(tail) - 1, -1, -1):
        acc += tail[i] * tail[i]
        edc[i] = acc
    level = lambda i: 10 * math.log10(max(edc[i], 1e-30) / edc[0])
    a = next((i for i in range(len(tail)) if level(i) <= -5), None)
    b = next((i for i in range(len(tail)) if level(i) <= -35), None)
    rt = 2 * (b - a) / rate if a is not None and b is not None else None
    return 10 * math.log10(energy / dry), rt


def check_reverb(folder, rate, channels, tol):
    import os
    want = {}
    with open(os.path.join(folder, "expected_reverb.txt")) as f:
        for line in f:
            name, decay, wet, effective = line.split()
            want[name] = (float(decay), float(wet), float(effective))
    got = {}
    print(f"{'scene':>8} {'decay of space':>16} {'RT60 expected':>14} {'measured RT60 s':>16} {'tail/click dB':>14}")
    for name in ("field", "gorge", "tunnel", "room"):
        level, rt = reverb_track(os.path.join(folder, name + ".raw"), rate, channels)
        got[name] = (level, rt)
        print(f"{name:>8} {want[name][0]:16.2f} {want[name][2]:14.2f} {('none' if rt is None else f'{rt:.2f}'):>16} {('none' if level is None else f'{level:.1f}'):>14}")
    ok = got["field"][1] is None
    print(f"  field has no tail: {'ok' if ok else 'FAIL'}")
    order = [got[n][1] for n in ("gorge", "tunnel", "room")]
    ordered = all(v is not None for v in order) and order[0] < order[1] < order[2]
    print(f"  gorge < tunnel < room: {'ok' if ordered else 'FAIL'}")
    close = True
    for n in ("tunnel", "room"):
        rt = got[n][1]
        good = rt is not None and abs(rt - want[n][2]) <= tol * want[n][2]
        close &= good
        print(f"  {n} within {tol:.0%} of the RT60 its sends make: {'ok' if good else 'FAIL'}")
    return ok and ordered and close


def db(a, b):
    return 10 * math.log10(max(a, 1e-20) / max(b, 1e-20))


def main():
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    p.add_argument("open", nargs="?")
    p.add_argument("wall", nargs="?")
    p.add_argument("--echo", help="track with one click: check the delays of its echoes")
    p.add_argument("--reverb-dir", help="folder written by export_reverb_tracks: check the decay of room, tunnel, gorge, field")
    p.add_argument("--rt-tol", type=float, default=0.35, help="how far a measured RT60 may be from the predicted, as a fraction")
    p.add_argument("--echo-dir", help="folder written by export_echo_tracks: check gorge, field and room")
    p.add_argument("--expect", default="none", help="expected delays in ms, comma separated, or none")
    p.add_argument("--tol", type=float, default=1.0, help="ms a peak may be off")
    p.add_argument("--rate", type=int, default=48000)
    p.add_argument("--channels", type=int, default=2)
    p.add_argument("--min-high-cut", type=float, default=15.0, help="dB the band above 5 kHz must drop")
    p.add_argument("--max-low-loss", type=float, default=8.0, help="dB the band below 300 Hz may drop")
    a = p.parse_args()
    if a.reverb_dir:
        return 0 if check_reverb(a.reverb_dir, a.rate, a.channels, a.rt_tol) else 1
    if a.echo_dir:
        import os
        ok = True
        with open(os.path.join(a.echo_dir, "expected.txt")) as f:
            for line in f:
                name, _, vals = line.strip().partition(" ")
                want = [float(v) for v in vals.split(",") if v]
                # The ms in the room are for the walls the geometry says; only the first peak is compared there.
                ok &= check_echo(os.path.join(a.echo_dir, name + ".raw"), want, a.tol, a.rate, a.channels)
        return 0 if ok else 1
    if a.echo:
        want = [] if a.expect == "none" else [float(v) for v in a.expect.split(",")]
        return 0 if check_echo(a.echo, want, a.tol, a.rate, a.channels) else 1
    if not a.open or not a.wall:
        p.error("give open.raw and wall.raw, or --echo / --echo-dir")
    o, w = bands(a.open, a.rate, a.channels), bands(a.wall, a.rate, a.channels)
    print(f"{'band, Hz':>10} {'open dB':>9} {'wall dB':>9} {'change dB':>10}")
    ref = max(o)
    for (name, _, _), x, y in zip(BANDS, o, w):
        print(f"{name:>10} {db(x, ref):9.1f} {db(y, ref):9.1f} {db(y, x):10.1f}")
    low, high = db(w[0], o[0]), db(w[-1], o[-1])
    ok = high <= -a.min_high_cut and low >= -a.max_low_loss
    print(f"highs {high:.1f} dB (need <= -{a.min_high_cut}), lows {low:.1f} dB (need >= -{a.max_low_loss}): {'ok' if ok else 'FAIL'}")
    return 0 if ok else 1


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