#!/usr/bin/env python3
"""Compare synthesized words against reference WAV directories.

Used to tune the Half-Life VOX preset against the game's word assets
(`suchipi/half-life-vox-console`, `public/vox/*.wav` - Valve assets, keep them
outside this repository).

Usage:
    vox_compare.py REFERENCE_DIR SYNTH_DIR [SYNTH_DIR ...]

Directories must hold the same word file names (one WAV per word). Prints, per
directory: median F0 (cepstrum), total audible duration, energy-weighted
spectral centroid and the >2 kHz energy ratio. These are the metrics used to
pick the VOX preset values in setup/veskslop-tts.sh.

Requires numpy.
"""
import os
import sys

import numpy as np
import wave


def frames(path):
    with wave.open(path) as w:
        data = w.readframes(w.getnframes())
        rate = w.getframerate()
    samples = np.frombuffer(data, dtype="<i2").astype(float)
    loud = np.nonzero(np.abs(samples) > 200)[0]
    if len(loud) == 0:
        return None, rate
    return samples[loud[0]:loud[-1] + 1], rate


def f0_cepstrum(seg, rate):
    win = min(2048, (len(seg) // 2) * 2)
    if win < 256:
        return 0.0
    values = []
    for i in range(0, max(1, len(seg) - win), win // 2):
        fr = seg[i:i + win]
        if (fr ** 2).sum() < 1e6:
            continue
        fr = fr - fr.mean()
        spectrum = np.abs(np.fft.rfft(fr * np.hanning(win))) + 1e-9
        cepstrum = np.fft.irfft(np.log(spectrum))
        lo, hi = int(rate / 400), int(rate / 60)
        if hi <= lo:
            continue
        lag = lo + int(np.argmax(cepstrum[lo:hi]))
        values.append(rate / lag)
    values = sorted(v for v in values if 50 < v < 400)
    return values[len(values) // 2] if values else 0.0


def brightness(seg, rate):
    win = min(2048, (len(seg) // 2) * 2)
    cent, high, energy = [], [], []
    for i in range(0, max(1, len(seg) - win), win // 2):
        fr = seg[i:i + win]
        if len(fr) < win:
            break
        e = (fr ** 2).sum() / win
        if e < 1e5:
            continue
        fr = fr - fr.mean()
        power = np.abs(np.fft.rfft(fr * np.hanning(win))) ** 2
        freq = np.fft.rfftfreq(win, 1 / rate)
        total = power.sum() + 1e-9
        cent.append((power * freq).sum() / total)
        high.append(power[(freq >= 2000) & (freq < 5000)].sum() / total)
        energy.append(e)
    if not cent:
        return 0.0, 0.0
    cent = np.array(cent)
    high = np.array(high)
    energy = np.array(energy)
    keep = energy >= np.percentile(energy, 70)
    return float(np.median(cent[keep])), float(np.median(high[keep]))


def summarize(directory, words):
    f0s, duration, cents, highs = [], 0.0, [], []
    for word in words:
        path = os.path.join(directory, word + ".wav")
        if not os.path.exists(path):
            continue
        seg, rate = frames(path)
        if seg is None or len(seg) == 0:
            continue
        duration += len(seg) / rate
        value = f0_cepstrum(seg, rate)
        if value > 0:
            f0s.append(value)
        c, h = brightness(seg, rate)
        cents.append(c)
        highs.append(h)
    f0s.sort()
    return {
        "n": len(f0s),
        "f0": f0s[len(f0s) // 2] if f0s else 0.0,
        "dur": duration,
        "cent": float(np.median(cents)) if cents else 0.0,
        "high": float(np.median(highs)) if highs else 0.0,
    }


def main():
    if len(sys.argv) < 3:
        print(__doc__)
        return 2
    reference = sys.argv[1]
    candidates = sorted(
        name[:-4] for name in os.listdir(reference)
        if name.endswith(".wav") and os.path.isfile(os.path.join(reference, name))
    )
    # Keep only words present in every directory so the summaries compare.
    words = [
        word for word in candidates
        if all(os.path.exists(os.path.join(d, word + ".wav")) for d in sys.argv[1:])
    ]
    for directory in sys.argv[1:]:
        s = summarize(directory, words)
        print(
            "%-28s words=%3d f0=%6.1fHz dur=%6.2fs cent=%5.0fHz high2k=%.4f"
            % (os.path.basename(directory.rstrip("/")), s["n"], s["f0"], s["dur"],
               s["cent"], s["high"])
        )
    return 0


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