"""Compare decoded audio streams using aligned float PCM arrays."""

import argparse
import json
import subprocess
from pathlib import Path

import numpy as np


def decode(path: Path) -> np.ndarray:
    raw = subprocess.check_output([
        "ffmpeg", "-hide_banner", "-loglevel", "error", "-i", str(path),
        "-map", "0:a:0", "-f", "f32le", "-acodec", "pcm_f32le", "-",
    ])
    return np.frombuffer(raw, dtype="<f4").reshape(-1, 2)


def metrics(reference: np.ndarray, candidate: np.ndarray) -> dict:
    count = min(len(reference), len(candidate))
    reference = reference[:count]
    candidate = candidate[:count]
    error = candidate - reference
    signal_power = np.maximum(np.mean(reference**2, axis=0), 1e-30)
    noise_power = np.maximum(np.mean(error**2, axis=0), 1e-30)
    first_250ms = error[:8000]
    first_signal_power = np.maximum(np.mean(reference[:8000] ** 2, axis=0), 1e-30)
    first_noise_power = np.maximum(np.mean(first_250ms**2, axis=0), 1e-30)
    return {
        "reference_samples": len(reference),
        "candidate_samples": len(candidate),
        "compared_samples": count,
        "snr_db_by_channel": (10.0 * np.log10(signal_power / noise_power)).tolist(),
        "rmse": float(np.sqrt(np.mean(error**2))),
        "first_250ms_snr_db_by_channel": (
            10.0 * np.log10(first_signal_power / first_noise_power)
        ).tolist(),
        "first_250ms_rmse": float(np.sqrt(np.mean(first_250ms**2))),
        "max_abs_error": float(np.max(np.abs(error))),
    }


parser = argparse.ArgumentParser()
parser.add_argument("--reference", type=Path, required=True)
parser.add_argument("--candidate", action="append", default=[], metavar="NAME=PATH")
parser.add_argument("--output", type=Path)
args = parser.parse_args()

reference = decode(args.reference)
report = {"reference": str(args.reference), "candidates": {}}
for value in args.candidate:
    if "=" not in value:
        raise ValueError(f"candidate must be NAME=PATH, got {value!r}")
    name, raw_path = value.split("=", 1)
    report["candidates"][name] = {
        "path": raw_path,
        **metrics(reference, decode(Path(raw_path))),
    }

serialized = json.dumps(report, indent=2)
if args.output is not None:
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(serialized + "\n", encoding="utf-8")
print(serialized)