"""Spectrogram features of the dataset clips, cached on disk.

Each clip becomes a log-magnitude spectrogram in dB relative to its own
maximum, floored at -80 dB and scaled to [0, 1], over 0 to FMAX Hz. The level
of the clip therefore does not enter (the WAVs are peak-normalized anyway).
The frequency resolution matters more than the time resolution here: the
sheet's low partials are tens of hertz apart and ring for seconds, so the
window is long (N_FFT / 48000 = 171 ms) and the hop coarse (21 ms).
"""

import json
import os
import sys
import wave
from concurrent.futures import ProcessPoolExecutor

import numpy as np

FS = 48000
N_FFT = 8192
HOP = 1024
FMAX = 6000.0
DB_FLOOR = 80.0
N_BINS = int(FMAX / FS * N_FFT)  # 1024 bins of 5.86 Hz
# Log-frequency variant: 48 bins per octave over 8 octaves from F_LO. On this
# axis a change of the sheet's frequency scale (kappa / L^2) translates the
# whole pattern of partials, which a convolutional network detects naturally.
F_LO = 20.0
BINS_PER_OCTAVE = 48
N_LOGF = 8 * BINS_PER_OCTAVE


def read_wav(path):
    with wave.open(path) as w:
        assert w.getnchannels() == 1 and w.getsampwidth() == 2 and w.getframerate() == FS, path
        x = np.frombuffer(w.readframes(w.getnframes()), dtype="<i2").astype(np.float32) / 32768.0
    return x


def spectrogram(x):
    """(N_BINS, n_frames) float16 in [0, 1]; 1 at the clip's loudest bin."""
    win = np.hanning(N_FFT + 1)[:-1].astype(np.float32)
    n_frames = 1 + (len(x) - N_FFT) // HOP
    frames = np.lib.stride_tricks.as_strided(x, shape=(n_frames, N_FFT), strides=(x.strides[0] * HOP, x.strides[0]))
    spec = np.abs(np.fft.rfft(frames * win, axis=1))[:, :N_BINS].T
    db = 20 * np.log10(spec + 1e-9)
    db -= db.max()
    return np.clip(db / DB_FLOOR + 1, 0, 1).astype(np.float16)


_FB = None


def logf_filterbank():
    """(N_LOGF, N_BINS) triangular filters centered on the log-spaced frequencies, at least one linear bin wide."""
    global _FB
    if _FB is None:
        f_lin = np.arange(N_BINS) * FS / N_FFT
        centers = F_LO * 2 ** (np.arange(N_LOGF) / BINS_PER_OCTAVE)
        fb = np.zeros((N_LOGF, N_BINS), np.float32)
        for k, fc in enumerate(centers):
            half = max(fc * (2 ** (1 / BINS_PER_OCTAVE) - 1), FS / N_FFT)
            w = np.clip(1 - np.abs(f_lin - fc) / half, 0, None)
            fb[k] = w / max(w.sum(), 1e-9)
        _FB = fb
    return _FB


def spectrogram_logf(x):
    """(N_LOGF, n_frames) float16 in [0, 1] on a log-frequency axis."""
    win = np.hanning(N_FFT + 1)[:-1].astype(np.float32)
    n_frames = 1 + (len(x) - N_FFT) // HOP
    frames = np.lib.stride_tricks.as_strided(x, shape=(n_frames, N_FFT), strides=(x.strides[0] * HOP, x.strides[0]))
    power = np.abs(np.fft.rfft(frames * win, axis=1))[:, :N_BINS] ** 2
    db = 10 * np.log10(logf_filterbank() @ power.T + 1e-18)
    db -= db.max()
    return np.clip(db / DB_FLOOR + 1, 0, 1).astype(np.float16)


def _one(args):
    clip_path, out_path, kind = args
    np.save(out_path, (spectrogram_logf if kind == "logf" else spectrogram)(read_wav(clip_path)))
    return out_path


def load_index(data_dir):
    rows = []
    with open(os.path.join(data_dir, "index.jsonl")) as f:
        for line in f:
            if line.strip():
                rows.append(json.loads(line))
    # Keep only clips whose WAV is present (the index may run ahead of a copy).
    return [r for r in rows if os.path.exists(os.path.join(data_dir, "clips", r["id"] + ".wav"))]


def build_cache(data_dir, cache_dir, workers=8, kind="lin"):
    """Computes the spectrograms (kind "lin" or "logf") of every indexed clip not yet cached; returns the index rows and the stacked features."""
    os.makedirs(cache_dir, exist_ok=True)
    rows = load_index(data_dir)
    todo = []
    for r in rows:
        out = os.path.join(cache_dir, r["id"] + ".npy")
        if not os.path.exists(out):
            todo.append((os.path.join(data_dir, "clips", r["id"] + ".wav"), out, kind))
    if todo:
        with ProcessPoolExecutor(workers) as ex:
            for _ in ex.map(_one, todo, chunksize=4):
                pass
    feats = np.stack([np.load(os.path.join(cache_dir, r["id"] + ".npy")) for r in rows])
    return rows, feats


if __name__ == "__main__":
    data_dir, cache_dir = sys.argv[1], sys.argv[2]
    rows, feats = build_cache(data_dir, cache_dir)
    print(len(rows), "clips, features", feats.shape, feats.dtype)