"""A similarity score for struck-sheet sounds that compares the partials with a
tolerance in frequency, and their levels and decays, without requiring the
lines to coincide.

Representation. Each clip is cut into K time slices (0-0.3, 0.3-1, 1-2, 2-3 s
by default). For each slice the mean power spectrum (8192-point windows,
5.9 Hz bins) is mapped onto a log-frequency axis (96 bins per octave, 30 Hz to
5 kHz) and smoothed there with a Gaussian whose width is the frequency
tolerance (2 percent by default), then converted to dB relative to the clip's
loudest value and floored. A line thus becomes a bump of width `tol`: two
lines a fraction of `tol` apart overlap almost completely, lines several `tol`
apart do not overlap at all, and in between the overlap falls off smoothly.
The level of a line is the height of its bump, and its decay is how the bump
shrinks from slice to slice, so the level, the position, and the decay of
every partial are all in the K x n_bins picture.

Score. The mean absolute difference in dB between the two pictures, over the
cells where either clip is above the floor (so silent regions do not dilute
it), with the later slices weighted equally to the first. Zero for identical
clips; smooth in the parameters of a simulation; insensitive to the level and
(through the log axis) treating a 2 percent shift at 100 Hz like one at 1 kHz.

    python3 ml/linesim.py TARGET.wav CLIP.wav [CLIP.wav ...] [--tol 0.02]
    python3 ml/linesim.py TARGET.wav --data DATASET_DIR --top 10     # rank a dataset
"""

import argparse
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from match import read_wav_any, onset, segment, FS  # noqa: E402

N_FFT = 8192
HOP = 1024
F_LO, F_HI = 30.0, 5000.0
BINS_PER_OCTAVE = 96
SLICES = ((0, 0.3), (0.3, 1.0), (1.0, 2.0), (2.0, 3.0))
FLOOR = -70.0

_fb_cache = {}


def logf_axis():
    n = int(np.log2(F_HI / F_LO) * BINS_PER_OCTAVE)
    return F_LO * 2 ** (np.arange(n) / BINS_PER_OCTAVE)


def smoothing_matrix(tol):
    """(n_log, n_lin): maps a linear power spectrum onto the log axis and smooths it with a Gaussian of relative width tol."""
    if tol in _fb_cache:
        return _fb_cache[tol]
    f_lin = np.arange(N_FFT // 2 + 1) * FS / N_FFT
    centers = logf_axis()
    W = np.exp(-0.5 * ((np.log(f_lin[None, 1:] / centers[:, None])) / tol) ** 2)
    W[:, :] /= np.maximum(W.sum(1, keepdims=True), 1e-12)
    W = np.concatenate([np.zeros((len(centers), 1)), W], 1).astype(np.float32)
    _fb_cache[tol] = W
    return W


def picture(seg, tol=0.02):
    """(K, n_log) dB picture of the clip; see the module docstring."""
    win = np.hanning(N_FFT + 1)[:-1].astype(np.float32)
    n_frames = 1 + (len(seg) - N_FFT) // HOP
    frames = np.lib.stride_tricks.as_strided(seg, shape=(n_frames, N_FFT), strides=(seg.strides[0] * HOP, seg.strides[0]))
    P = np.abs(np.fft.rfft(frames * win, axis=1)) ** 2
    fps = FS / HOP
    W = smoothing_matrix(tol)
    out = []
    for a, b in SLICES:
        sl = P[int(a * fps): max(int(a * fps) + 1, int(b * fps))].mean(0)
        out.append(W @ sl)
    db = 10 * np.log10(np.array(out) + 1e-20)
    db -= db.max()
    return np.maximum(db, FLOOR)


def distance(A, B):
    """Mean |dB difference| over the cells where either picture is above the floor, averaged over slices."""
    per_slice = []
    for a, b in zip(A, B):
        m = (a > FLOOR + 1) | (b > FLOOR + 1)
        per_slice.append(float(np.abs(a - b)[m].mean()) if m.any() else 0.0)
    return float(np.mean(per_slice)), per_slice


def load(path, align):
    x = read_wav_any(path)
    return segment(x, onset(x) if align else 0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("target")
    ap.add_argument("clips", nargs="*")
    ap.add_argument("--data", default=None)
    ap.add_argument("--top", type=int, default=10)
    ap.add_argument("--tol", type=float, default=0.02)
    ap.add_argument("--no-align", action="store_true", help="do not detect the onset of the target (simulated clips start at the impact)")
    args = ap.parse_args()
    T = picture(load(args.target, not args.no_align), args.tol)
    if args.clips:
        for c in args.clips:
            d, parts = distance(T, picture(load(c, False), args.tol))
            print(f"{os.path.basename(c):>16}: {d:6.2f} dB   slices " + " ".join(f"{p:.2f}" for p in parts))
    if args.data:
        rows = {json.loads(l)["id"]: json.loads(l) for l in open(os.path.join(args.data, "index.jsonl")) if l.strip()}
        ids = sorted(i for i in rows if os.path.exists(os.path.join(args.data, "clips", i + ".wav")))
        cache = os.path.join(args.data, "ml", f"linesim-tol{args.tol:g}.npz")
        if os.path.exists(cache) and list(np.load(cache)["ids"]) == ids:
            pics = np.load(cache)["pics"]
        else:
            pics = np.stack([picture(load(os.path.join(args.data, "clips", i + ".wav"), False), args.tol) for i in ids]).astype(np.float32)
            np.savez(cache, ids=np.array(ids), pics=pics)
        d = np.array([distance(T, p)[0] for p in pics])
        order = np.argsort(d)
        print(f"{len(ids)} clips; median distance {np.median(d):.2f} dB")
        for k in order[: args.top]:
            r = rows[ids[k]]
            print(f"{ids[k]} {d[k]:6.2f} dB | E {r['E'] / 1e9:.0f} rho {r['rho']:.0f} h {r['h'] * 1e3:.2f} {r['Lx'] / 0.0254:.1f}x{r['Ly'] / 0.0254:.1f} in v {r['strikeVelocity']:.1f} sig0 {r['sigma0']:.2f} sig1 {r['sigma1'] * 1e4:.1f} p2p {r['peakToPeak'] * 1e3:.1f}")
        return [ids[k] for k in order[: args.top]]


if __name__ == "__main__":
    main()