"""Measurement layer.

Pure numpy/scipy. No librosa — import cost and CPU headroom matter on a
2-vCPU Space where this runs several times a second on a live stream.

Everything is measured at 48 kHz. Input at any rate is resampled once on
the way in, so the K-weighting biquads (which are rate-specific) stay valid
and every metric is comparable across sources.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field, asdict
from typing import Optional

import numpy as np
from scipy.signal import lfilter, resample_poly, get_window

SR = 48_000

# Band edges in Hz. Seven bands, chosen to match how engineers actually talk
# about a mix rather than to divide the spectrum evenly.
BANDS: dict[str, tuple[float, float]] = {
    "sub": (20.0, 60.0),
    "bass": (60.0, 120.0),
    "lowmid": (120.0, 350.0),
    "mid": (350.0, 1500.0),
    "himid": (1500.0, 4000.0),
    "presence": (4000.0, 8000.0),
    "air": (8000.0, 16000.0),
}
BAND_ORDER = list(BANDS.keys())

BAND_LABELS = {
    "sub": "Sub 20–60",
    "bass": "Bass 60–120",
    "lowmid": "Low mid 120–350",
    "mid": "Mid 350–1.5k",
    "himid": "Hi mid 1.5–4k",
    "presence": "Presence 4–8k",
    "air": "Air 8–16k",
}

# ITU-R BS.1770-4 K-weighting, 48 kHz.
_K1_B = np.array([1.53512485958697, -2.69169618940638, 1.19839281085285])
_K1_A = np.array([1.0, -1.69065929318241, 0.73248077421585])
_K2_B = np.array([1.0, -2.0, 1.0])
_K2_A = np.array([1.0, -1.99004745483398, 0.99007225036621])

_KRUMHANSL_MAJOR = np.array(
    [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]
)
_KRUMHANSL_MINOR = np.array(
    [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17]
)
_NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]

_EPS = 1e-12


def _db(x: float) -> float:
    return 20.0 * math.log10(max(float(x), _EPS))


def _pdb(p: float) -> float:
    """Power (already squared) to dB."""
    return 10.0 * math.log10(max(float(p), _EPS))


# --------------------------------------------------------------------------
# input conditioning
# --------------------------------------------------------------------------

def to_float_stereo(sr: int, data: np.ndarray) -> np.ndarray:
    """Normalise any Gradio audio payload to float32 (n, 2) at 48 kHz."""
    x = np.asarray(data)
    if x.dtype.kind in "iu":
        info = np.iinfo(x.dtype)
        x = x.astype(np.float32) / max(abs(info.min), info.max)
    else:
        x = x.astype(np.float32, copy=False)

    if x.ndim == 1:
        x = x[:, None]
    if x.shape[1] > 2:  # some capture paths hand back (channels, n)
        if x.shape[0] <= 2:
            x = x.T
        else:
            x = x[:, :2]
    if x.shape[1] == 1:
        x = np.repeat(x, 2, axis=1)

    if sr != SR and x.shape[0] > 0:
        g = math.gcd(int(sr), SR)
        x = resample_poly(x, SR // g, int(sr) // g, axis=0).astype(np.float32)

    return np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)


# --------------------------------------------------------------------------
# loudness
# --------------------------------------------------------------------------

def _k_weight(x: np.ndarray) -> np.ndarray:
    y = lfilter(_K1_B, _K1_A, x, axis=0)
    return lfilter(_K2_B, _K2_A, y, axis=0)


def _block_loudness(xk: np.ndarray, win: int, hop: int) -> np.ndarray:
    """Per-block BS.1770 loudness in LKFS for a K-weighted signal."""
    n = xk.shape[0]
    if n < win:
        if n == 0:
            return np.array([])
        mean_sq = np.mean(xk**2, axis=0).sum()
        return np.array([-0.691 + _pdb(mean_sq)])
    starts = np.arange(0, n - win + 1, hop)
    out = np.empty(len(starts), dtype=np.float64)
    for i, s in enumerate(starts):
        blk = xk[s : s + win]
        out[i] = -0.691 + _pdb(np.mean(blk**2, axis=0).sum())
    return out


def loudness(x: np.ndarray) -> dict:
    """Integrated / short-term / range loudness, gated per BS.1770-4."""
    xk = _k_weight(x)
    blocks = _block_loudness(xk, int(0.400 * SR), int(0.100 * SR))
    short = _block_loudness(xk, int(3.0 * SR), int(1.0 * SR))

    lufs_i = float("-inf")
    if blocks.size:
        above_abs = blocks[blocks > -70.0]
        if above_abs.size:
            mean_pow = np.mean(10 ** (above_abs / 10.0))
            rel_gate = 10.0 * math.log10(max(mean_pow, _EPS)) - 10.0
            kept = above_abs[above_abs > rel_gate]
            pool = kept if kept.size else above_abs
            lufs_i = float(10.0 * math.log10(max(np.mean(10 ** (pool / 10.0)), _EPS)))

    lra = 0.0
    if short.size >= 3:
        valid = short[short > -70.0]
        if valid.size >= 3:
            lra = float(np.percentile(valid, 95) - np.percentile(valid, 10))

    return {
        "lufs_i": lufs_i,
        "lufs_s": float(short[-1]) if short.size else float("-inf"),
        "lra": lra,
        "short_term": short,
    }


def true_peak_db(x: np.ndarray) -> float:
    """dBTP via 4x oversampling (BS.1770 minimum)."""
    if x.shape[0] < 8:
        return _db(np.max(np.abs(x)) if x.size else 0.0)
    up = resample_poly(x, 4, 1, axis=0)
    return _db(float(np.max(np.abs(up))))


# --------------------------------------------------------------------------
# spectrum
# --------------------------------------------------------------------------

def spectrum(mono: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Averaged power spectrum via Welch-style overlapping Hann frames."""
    n = mono.shape[0]
    nfft = 8192 if n >= 8192 else 1 << max(8, int(math.log2(max(n, 256))))
    if n < nfft:
        mono = np.pad(mono, (0, nfft - n))
        n = nfft
    win = get_window("hann", nfft, fftbins=True)
    hop = nfft // 2
    acc = np.zeros(nfft // 2 + 1)
    count = 0
    for s in range(0, n - nfft + 1, hop):
        frame = mono[s : s + nfft] * win
        acc += np.abs(np.fft.rfft(frame)) ** 2
        count += 1
    if count:
        acc /= count
    freqs = np.fft.rfftfreq(nfft, 1.0 / SR)
    return freqs, acc


def band_energies(freqs: np.ndarray, power: np.ndarray) -> dict[str, float]:
    """Band power as dB relative to total 20 Hz–20 kHz power."""
    full = (freqs >= 20.0) & (freqs <= 20000.0)
    total = float(power[full].sum())
    out = {}
    for name, (lo, hi) in BANDS.items():
        sel = (freqs >= lo) & (freqs < hi)
        out[name] = _pdb(float(power[sel].sum()) / max(total, _EPS))
    return out


def spectral_shape(freqs: np.ndarray, power: np.ndarray) -> dict:
    sel = (freqs >= 40.0) & (freqs <= 16000.0)
    f, p = freqs[sel], power[sel]
    if not f.size or p.sum() <= _EPS:
        return {"centroid": 0.0, "tilt": 0.0, "rolloff85": 0.0, "flatness": 0.0}

    centroid = float((f * p).sum() / p.sum())

    # dB/octave tilt: least-squares fit of level against log2(frequency).
    logf = np.log2(f)
    logp = 10.0 * np.log10(np.maximum(p, _EPS))
    tilt = float(np.polyfit(logf, logp, 1)[0])

    csum = np.cumsum(p)
    rolloff = float(f[min(int(np.searchsorted(csum, 0.85 * csum[-1])), f.size - 1)])

    gmean = float(np.exp(np.mean(np.log(np.maximum(p, _EPS)))))
    flatness = gmean / float(np.mean(p) + _EPS)

    return {
        "centroid": centroid,
        "tilt": tilt,
        "rolloff85": rolloff,
        "flatness": float(flatness),
    }


# --------------------------------------------------------------------------
# stereo
# --------------------------------------------------------------------------

def stereo_image(x: np.ndarray) -> dict:
    left, right = x[:, 0], x[:, 1]
    denom = math.sqrt(float(np.mean(left**2)) * float(np.mean(right**2))) + _EPS
    corr = float(np.mean(left * right) / denom)

    mid = (left + right) * 0.5
    side = (left - right) * 0.5
    mid_p = float(np.mean(mid**2))
    side_p = float(np.mean(side**2))
    width = _pdb(side_p / max(mid_p, _EPS))

    # Mono-fold penalty: how much level is lost summing to mono.
    stereo_p = float(np.mean(x**2))
    mono_loss = _pdb(mid_p / max(stereo_p, _EPS))

    # Low-frequency correlation is the one that actually costs you on a
    # club system, so it gets measured separately.
    sub_corr = corr
    if x.shape[0] > 4096:
        nfft = 4096
        w = get_window("hann", nfft)
        bins = np.fft.rfftfreq(nfft, 1.0 / SR)
        m = (bins >= 20) & (bins < 120)
        limit = min(x.shape[0] - nfft, 24 * nfft)
        acc_l = acc_r = acc_lr = 0.0
        for s in range(0, max(limit, 1), nfft // 2):
            fl = np.fft.rfft(left[s : s + nfft] * w)
            fr = np.fft.rfft(right[s : s + nfft] * w)
            acc_l += float(np.sum(np.abs(fl[m]) ** 2))
            acc_r += float(np.sum(np.abs(fr[m]) ** 2))
            acc_lr += float(np.real(np.sum(fl[m] * np.conj(fr[m]))))
        d = math.sqrt(acc_l * acc_r) + _EPS
        sub_corr = float(acc_lr / d)

    return {
        "correlation": corr,
        "sub_correlation": sub_corr,
        "width_db": width,
        "mono_loss_db": mono_loss,
    }


# --------------------------------------------------------------------------
# rhythm + pitch
# --------------------------------------------------------------------------

def onset_envelope(mono: np.ndarray) -> tuple[np.ndarray, float]:
    nfft, hop = 2048, 512
    if mono.shape[0] < nfft * 4:
        return np.zeros(0), SR / hop
    win = get_window("hann", nfft)
    n_frames = 1 + (mono.shape[0] - nfft) // hop
    mags = np.empty((n_frames, nfft // 2 + 1), dtype=np.float32)
    for i in range(n_frames):
        s = i * hop
        mags[i] = np.abs(np.fft.rfft(mono[s : s + nfft] * win))
    logm = np.log1p(mags * 100.0)
    flux = np.maximum(np.diff(logm, axis=0), 0.0).sum(axis=1)
    if flux.size and flux.max() > 0:
        flux = flux / flux.max()
    return flux, SR / hop


def tempo_from_onsets(flux: np.ndarray, fps: float) -> dict:
    if flux.size < 64:
        return {"bpm": 0.0, "confidence": 0.0, "onset_rate": 0.0}

    env = flux - flux.mean()
    ac = np.correlate(env, env, mode="full")[env.size - 1 :]
    if ac[0] > 0:
        ac = ac / ac[0]

    lag_min = max(int(fps * 60.0 / 200.0), 2)
    lag_max = min(int(fps * 60.0 / 60.0), ac.size - 1)
    if lag_max <= lag_min:
        return {"bpm": 0.0, "confidence": 0.0, "onset_rate": 0.0}

    window = ac[lag_min:lag_max]
    best = int(np.argmax(window)) + lag_min
    conf = float(max(window.max(), 0.0))
    bpm = 60.0 * fps / best

    # Octave correction — autocorrelation happily locks onto half or double.
    while bpm < 70.0:
        bpm *= 2.0
    while bpm > 190.0:
        bpm /= 2.0

    thresh = flux.mean() + flux.std()
    peaks = (flux[1:-1] > thresh) & (flux[1:-1] > flux[:-2]) & (flux[1:-1] >= flux[2:])
    onset_rate = int(np.sum(peaks)) / (flux.size / fps) if flux.size else 0.0

    return {"bpm": float(bpm), "confidence": conf, "onset_rate": float(onset_rate)}


def key_estimate(freqs: np.ndarray, power: np.ndarray) -> dict:
    sel = (freqs >= 55.0) & (freqs <= 2200.0)
    f, p = freqs[sel], power[sel]
    if not f.size or p.sum() <= _EPS:
        return {"key": "—", "confidence": 0.0}

    midi = 69.0 + 12.0 * np.log2(f / 440.0)
    pc = np.mod(np.round(midi).astype(int), 12)
    chroma = np.zeros(12)
    np.add.at(chroma, pc, np.sqrt(p))
    if chroma.sum() <= _EPS:
        return {"key": "—", "confidence": 0.0}
    chroma = chroma / chroma.sum()

    scored: list[tuple[float, str]] = []
    for root in range(12):
        rotated = np.roll(chroma, -root)
        for profile, quality in ((_KRUMHANSL_MAJOR, ""), (_KRUMHANSL_MINOR, "m")):
            prof = profile / profile.sum()
            if np.std(rotated) < _EPS:
                continue
            score = float(np.corrcoef(rotated, prof)[0, 1])
            scored.append((score, f"{_NOTE_NAMES[root]}{quality}"))

    if not scored:
        return {"key": "—", "confidence": 0.0}
    scored.sort(reverse=True)
    margin = scored[0][0] - (scored[1][0] if len(scored) > 1 else 0.0)
    return {"key": scored[0][1], "confidence": float(max(0.0, margin))}


# --------------------------------------------------------------------------
# top-level report
# --------------------------------------------------------------------------

@dataclass
class Report:
    duration: float = 0.0
    lufs_i: float = float("-inf")
    lufs_s: float = float("-inf")
    lra: float = 0.0
    true_peak: float = -120.0
    sample_peak: float = -120.0
    rms: float = -120.0
    crest: float = 0.0
    psr: float = 0.0
    bands: dict = field(default_factory=dict)
    ratios: dict = field(default_factory=dict)
    shape: dict = field(default_factory=dict)
    stereo: dict = field(default_factory=dict)
    rhythm: dict = field(default_factory=dict)
    key: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        d = asdict(self)
        for k, v in list(d.items()):
            if isinstance(v, float) and math.isinf(v):
                d[k] = -120.0
        return d


def analyze(sr: int, data: np.ndarray, *, fast: bool = False) -> Optional[Report]:
    """Full measurement pass. `fast=True` skips tempo/key for the live loop."""
    x = to_float_stereo(sr, data)
    if x.shape[0] < SR // 20:
        return None

    mono = x.mean(axis=1)
    rms = float(np.sqrt(np.mean(mono**2)))
    sample_peak = float(np.max(np.abs(x)))

    loud = loudness(x)
    freqs, power = spectrum(mono)
    bands = band_energies(freqs, power)

    rep = Report(
        duration=x.shape[0] / SR,
        lufs_i=loud["lufs_i"],
        lufs_s=loud["lufs_s"],
        lra=loud["lra"],
        true_peak=true_peak_db(x) if not fast else _db(sample_peak),
        sample_peak=_db(sample_peak),
        rms=_db(rms),
        crest=_db(sample_peak) - _db(rms),
        bands=bands,
        shape=spectral_shape(freqs, power),
        stereo=stereo_image(x),
    )

    # Peak-to-short-term-loudness ratio: the "is it still punchy" number.
    if loud["short_term"].size and math.isfinite(rep.true_peak):
        rep.psr = float(rep.true_peak - float(loud["short_term"][-1]))

    rep.ratios = {
        "sub_vs_bass": bands["sub"] - bands["bass"],
        "mud": bands["lowmid"] - bands["mid"],
        "harsh": bands["himid"] - bands["mid"],
        "air": bands["air"] - bands["mid"],
        "tilt_low_high": (bands["sub"] + bands["bass"]) - (bands["presence"] + bands["air"]),
    }

    if not fast:
        flux, fps = onset_envelope(mono)
        rep.rhythm = tempo_from_onsets(flux, fps)
        rep.key = key_estimate(freqs, power)
    else:
        rep.rhythm = {"bpm": 0.0, "confidence": 0.0, "onset_rate": 0.0}
        rep.key = {"key": "—", "confidence": 0.0}

    return rep