#!/usr/bin/env python3
# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""Write the audio front-end reference for `eg2-cli audiocheck`: for every clip of tests/audio/, the Python front end's
16 kHz waveform and log-mel features (raw float32 files in tests/audio/ref/) and tests/audio-preprocess-cases.json
with the expected counts and the tolerances the Swift front end must meet.

    python3 make_audio_preprocess_cases.py [--audio audio] [--out audio-preprocess-cases.json]

Tolerances: lossless files decode identically on both sides and the resampler is the same double-precision code, so
the waveform must match to 1e-6 (0 for files already at 16 kHz); AAC files decode with the same Apple decoder (1e-6);
MP3 files are decoded by libsndfile here (when installed) and by Apple's decoder in Swift, which differ by up to ~3e-5
per sample, which in near-silent frames moves ln(mel + 0.001) by up to ~1 (tolerance 1.0); for the other formats the
features must match to 0.02 (the single-precision FFTs of numpy and Accelerate round differently, ~0.005 worst).
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "python"))
import eg2_audio_frontend as fe  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--audio", default=str(HERE / "audio"))
    ap.add_argument("--out", default=str(HERE / "audio-preprocess-cases.json"))
    args = ap.parse_args()
    audio = Path(args.audio)
    ref = audio / "ref"
    ref.mkdir(exist_ok=True)
    manifest = json.loads((audio / "manifest.json").read_text(encoding="utf-8"))
    cases = []
    for it in manifest:
        d = fe.decode(audio / it["file"])
        wav = fe.resample_16k(fe.to_mono(d.samples), d.sample_rate)
        inputs = fe.prepare_audio_inputs(wav)
        feats = fe.log_mel(wav)
        wave_file = f"ref/{it['name']}.wave16k.f32"
        feat_file = f"ref/{it['name']}.features.f32"
        wav.astype("<f4").tofile(audio / wave_file)
        feats.astype("<f4").tofile(audio / feat_file)
        is_mp3 = it["file"].endswith(".mp3")
        cases.append(dict(file=f"audio/{it['file']}", sample_rate=d.sample_rate, channels=int(d.samples.shape[1]), n_samples16k=int(len(wav)),
                          n_frames=inputs.n_frames, n_tokens=inputs.n_tokens, bucket=inputs.bucket,
                          wave_file=f"audio/{wave_file}", features_file=f"audio/{feat_file}",
                          wave_tolerance=1e-4 if is_mp3 else (0.0 if d.sample_rate == 16000 and not it["lossy"] else 1e-6),
                          feature_tolerance=1.0 if is_mp3 else 0.02, python_decoder=d.decoder,
                          note=("mp3: libsndfile (Python) vs Apple (Swift) decoders differ by ~3e-5 per sample" if is_mp3 else None)))
        print(cases[-1], flush=True)
    Path(args.out).write_text(json.dumps(cases, indent=1) + "\n", encoding="utf-8")
    print(f"{len(cases)} cases -> {args.out}")


if __name__ == "__main__":
    main()