# SPDX-License-Identifier: Apache-2.0
"""Dump the Python metrics as the reference for the browser port (web/dsp.js).

Each audio file is loaded exactly as the desktop app loads it
(`master_core.AudioFile`); the mono signal the metrics see is written out next
to every metric's `compute()` result. `parity.mjs` feeds that same signal to
dsp.js, so the comparison isolates the DSP port from decoder differences.

  uv run --no-project --with librosa --with pyloudnorm --with scipy --with mutagen \
      python tests/parity/reference.py OUT_DIR AUDIO...

Writes OUT_DIR/<file name>/{mono.f32, reference.json, spectrogram.f32}.
"""

import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from master_core import AudioFile  # noqa: E402
from metrics import METRICS  # noqa: E402


def _jsonable(value):
  if isinstance(value, np.ndarray):
    return [float(v) if np.isfinite(v) else None for v in value.astype(np.float64).ravel()]
  if isinstance(value, (float, np.floating)):
    return float(value) if np.isfinite(value) else None
  return value


def main() -> None:
  out = Path(sys.argv[1])
  out.mkdir(parents=True, exist_ok=True)
  for path in sys.argv[2:]:
    name = Path(path).name
    try:
      audio = AudioFile(path)
    except Exception as exc:  # librosa can't read every container the browser can
      print(f"{name}: skipped ({exc})")
      continue

    target = out / name
    target.mkdir(exist_ok=True)
    audio.y_mono.astype(np.float32).tofile(target / "mono.f32")

    results = {}
    for metric_id, metric in METRICS.items():
      data = metric.compute(audio)
      if metric_id == "spectrogram":
        s_db = np.ascontiguousarray(data["s_db"], dtype=np.float32)
        s_db.tofile(target / "spectrogram.f32")
        results[metric_id] = {
          "rows": int(s_db.shape[0]), "cols": int(s_db.shape[1]),
          "times": _jsonable(data["times"]), "freqs": _jsonable(data["freqs"]),
          "nyquist": float(data["nyquist"]),
        }
      else:
        results[metric_id] = {key: _jsonable(val) for key, val in data.items()}

    meta = {
      "file": name, "sr": int(audio.sr), "n": int(len(audio.y_mono)),
      "song_name": audio.song_name,
      "max_amplitude": float(audio.max_amplitude), "avg_amplitude": float(audio.avg_amplitude),
    }
    (target / "reference.json").write_text(json.dumps({"meta": meta, "metrics": results}))
    print(f"{name}: {audio.sr} Hz, {len(audio.y_mono) / audio.sr:.2f} s, song_name={audio.song_name!r}")


if __name__ == "__main__":
  main()