"""Error measures between a solution and a reference, as defined in each challenge README."""
from __future__ import annotations

import numpy as np


def waveform_l2(a: np.ndarray, r: np.ndarray) -> float:
    return float(np.linalg.norm(a - r) / np.linalg.norm(r))


def spectrogram(x: np.ndarray, n_fft: int, hop: int) -> np.ndarray:
    """Magnitude STFT with a periodic Hann window, no padding. Shape (frames, n_fft//2 + 1)."""
    j = np.arange(n_fft)
    g = 0.5 * (1.0 - np.cos(2.0 * np.pi * j / n_fft))
    m = (len(x) - n_fft) // hop + 1
    idx = j[None, :] + hop * np.arange(m)[:, None]
    frames = x[idx] * g[None, :]
    return np.abs(np.fft.rfft(frames, axis=1))


def spectrogram_l2(a: np.ndarray, r: np.ndarray, n_fft: int, hop: int) -> float:
    sa = spectrogram(a, n_fft, hop)
    sr = spectrogram(r, n_fft, hop)
    return float(np.linalg.norm(sa - sr) / np.linalg.norm(sr))


def compute_metric(name: str, a: np.ndarray, r: np.ndarray, params: dict) -> float:
    if name == "waveform_l2":
        return waveform_l2(a, r)
    if name == "spectrogram_l2":
        return spectrogram_l2(a, r, int(params["n_fft"]), int(params["hop"]))
    raise KeyError(f"unknown metric {name}")


def metric_names(challenge: dict) -> list[str]:
    return [k for k in challenge["metrics"] if k != "primary"]