"""Quantization, audio IO, objective metrics and plots for the FocalCodec lab.

Kept free of Gradio / torch.hub so it can be exercised standalone.
"""

from __future__ import annotations

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
import soundfile as sf
import torch
from pystoi import stoi
from scipy.signal import resample_poly, stft

QUANT_LEVELS = ("int4", "int8", "int16", "int32", "fp4", "fp8", "fp16", "fp32")

TOKEN_RATE_HZ = 12.5
BITS_PER_TOKEN = 13
BITRATE_KBPS = TOKEN_RATE_HZ * BITS_PER_TOKEN / 1000.0

_BITS = {"int4": 4, "int8": 8, "int16": 16, "int32": 32, "fp4": 4, "fp8": 8, "fp16": 16, "fp32": 32}
_FP_FORMATS = {"fp4": (2, 1), "fp8": (4, 3)}  # (exponent bits, mantissa bits) -> E2M1 / E4M3
_FP_MAX_NORM = {"fp4": 6.0, "fp8": 448.0}  # E2M1 max; E4M3-FN finite max (top encoding is NaN)


def _ref_scale(w: torch.Tensor) -> torch.Tensor:
    """Per-output-channel absmax, per-tensor for 1-D params."""
    if w.dim() > 1:
        return w.abs().amax(dim=tuple(range(1, w.dim())), keepdim=True)
    return w.abs().amax()


def int_scale(w: torch.Tensor, bits: int) -> torch.Tensor:
    return (_ref_scale(w) / (1 << (bits - 1))).clamp_min(1e-12)


def int_levels(w: torch.Tensor, bits: int) -> torch.Tensor:
    """Signed integer levels before rescaling (what an int-format export stores)."""
    qmax = 1 << (bits - 1)
    return (w / int_scale(w, bits)).round().clamp(-qmax, qmax - 1)


def fp_scale(w: torch.Tensor, level: str) -> torch.Tensor:
    return (_ref_scale(w) / _FP_MAX_NORM[level]).clamp_min(1e-12)


def fp_levels(w: torch.Tensor, level: str) -> torch.Tensor:
    """Values on the format's own grid (what a float-format export stores, unscaled)."""
    e_bits, m_bits = _FP_FORMATS[level]
    bias = (1 << (e_bits - 1)) - 1
    max_exp = (1 << e_bits) - 1 - bias
    x = w / fp_scale(w, level)
    ax = x.abs()
    exp = torch.floor(torch.log2(ax.clamp_min(1e-30))).clamp(1 - bias, max_exp)
    step = torch.pow(2.0, exp - m_bits)
    return torch.sign(x) * ((ax / step).round() * step)


def _quant_int(w: torch.Tensor, bits: int) -> torch.Tensor:
    if bits >= 32:
        return w
    return int_levels(w, bits) * int_scale(w, bits)


def _quant_fp(w: torch.Tensor, level: str) -> torch.Tensor:
    return fp_levels(w, level) * fp_scale(w, level)


def quantize_state(state: dict, level: str) -> dict:
    """Fake-quantize every float tensor in `state` to `level` (per-channel absmax scaling)."""
    if level == "fp32":
        return dict(state)
    out = {}
    for key, value in state.items():
        if not value.is_floating_point():
            out[key] = value
        elif level == "fp16":
            out[key] = value.half().float()
        elif level.startswith("int"):
            out[key] = _quant_int(value, int(level[3:]))
        else:
            out[key] = _quant_fp(value, level)
    return out


def model_size_mb(state: dict, level: str) -> float:
    """Stored size of `state` at `level`; non-float buffers keep their own width."""
    return sum(
        v.numel() * (v.element_size() if not v.is_floating_point() else _BITS[level] / 8)
        for v in state.values()
    ) / 1e6


def load_audio(path: str, target_sr: int, max_seconds: float | None = None) -> tuple[np.ndarray, int]:
    """Read `path` as mono float32, resample to `target_sr`, optionally truncate."""
    x, sr = sf.read(path, dtype="float32", always_2d=True)
    x = x.mean(axis=1)
    if sr != target_sr:
        g = np.gcd(sr, target_sr)
        x = resample_poly(x, target_sr // g, sr // g).astype(np.float32)
        sr = target_sr
    if max_seconds is not None and len(x) > int(max_seconds * sr):
        x = x[: int(max_seconds * sr)]
    return x, sr


def si_sdr(ref: np.ndarray, rec: np.ndarray) -> float:
    ref = ref - ref.mean()
    rec = rec - rec.mean()
    denom = float(ref @ ref) + 1e-12
    target = (float(rec @ ref) / denom) * ref
    noise = rec - target
    return float(10 * np.log10((float(target @ target) + 1e-12) / (float(noise @ noise) + 1e-12)))


def log_spectral_distance(ref: np.ndarray, rec: np.ndarray, sr: int) -> float:
    """Frame-averaged log-spectral distance (dB) between reference and reconstruction."""
    _, _, R = stft(ref, fs=sr, nperseg=1024, noverlap=768)
    _, _, S = stft(rec, fs=sr, nperseg=1024, noverlap=768)
    diff = 20 * np.log10(np.abs(R) + 1e-8) - 20 * np.log10(np.abs(S) + 1e-8)
    return float(np.mean(np.sqrt(np.mean(diff**2, axis=0))))


def evaluate(
    ref: np.ndarray, rec: np.ndarray, sr: int, rtf: float, size_mb: float, level: str,
    bitrate_kbps: float | None = None,
) -> dict:
    n = min(len(ref), len(rec))
    ref, rec = ref[:n], rec[:n]
    if bitrate_kbps is None:
        bitrate_kbps = BITRATE_KBPS
    return {
        "Level": level,
        "SI-SDR (dB)": round(si_sdr(ref, rec), 2),
        "STOI": round(float(stoi(ref, rec, sr, extended=False)), 3),
        "LSD (dB)": round(log_spectral_distance(ref, rec, sr), 2),
        "Model size (MB)": round(size_mb, 1),
        "Bitrate (kbps)": round(bitrate_kbps, 4),
        "RTF": round(rtf, 3),
    }


def spectrogram_figure(ref: np.ndarray, rec: np.ndarray, sr: int, title: str):
    fig, axes = plt.subplots(1, 2, figsize=(9.0, 2.6), sharey=True)
    for ax, sig, label in zip(axes, (ref, rec), ("Input", "Reconstruction")):
        ax.specgram(sig, NFFT=1024, noverlap=768, Fs=sr, cmap="magma")
        ax.set_title(label, fontsize=9)
        ax.set_xlabel("Time (s)")
    axes[0].set_ylabel("Frequency (Hz)")
    fig.suptitle(title, fontsize=10)
    fig.tight_layout()
    return fig


def metrics_figure(rows: list[dict]):
    levels = [r["Level"] for r in rows]
    panels = (("SI-SDR (dB)", "#4c72b0"), ("STOI", "#55a868"), ("Model size (MB)", "#c44e52"))
    fig, axes = plt.subplots(1, len(panels), figsize=(11.5, 2.9))
    for ax, (key, color) in zip(axes, panels):
        ax.bar(levels, [r[key] for r in rows], color=color)
        ax.set_title(key, fontsize=9)
        ax.tick_params(axis="x", rotation=45)
        ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    return fig


if __name__ == "__main__":
    torch.manual_seed(0)
    w = (torch.randn(64, 32) * 0.05).requires_grad_(False)
    for level in QUANT_LEVELS:
        q = quantize_state({"w": w}, level)["w"]
        assert torch.isfinite(q).all(), f"{level} not finite"
        assert q.abs().max() <= w.abs().max() * 1.01 + 1e-9, f"{level} broke the magnitude bound"
    # the packed levels a real export stores must reproduce the fake-quant demo exactly
    assert torch.equal(int_levels(w, 8) * int_scale(w, 8), quantize_state({"w": w}, "int8")["w"])
    assert torch.equal(fp_levels(w, "fp8") * fp_scale(w, "fp8"), quantize_state({"w": w}, "fp8")["w"])
    # fp8 grid values are exactly float8_e4m3fn-representable, so no precision is lost on store
    grid8 = fp_levels(w, "fp8")
    assert torch.equal(grid8.to(torch.float8_e4m3fn).float(), grid8), "fp8 grid escapes E4M3-FN"
    grid4 = fp_levels(w, "fp4")
    assert grid4.abs().max() <= _FP_MAX_NORM["fp4"] * 1.0000001, "fp4 grid escapes E2M1"
    assert model_size_mb({"w": w}, "fp32") == w.numel() * 4 / 1e6
    assert model_size_mb({"i": torch.zeros(8, dtype=torch.int64)}, "int4") == 8 * 8 / 1e6  # buffers keep 8 bytes
    x = np.random.randn(16000).astype(np.float32)
    assert si_sdr(x, x) > 60, "SI-SDR of identical signals should be large"
    assert si_sdr(x, 3 * x) > 60, "SI-SDR must be scale invariant"
    assert abs(evaluate(x, x, 16000, 1.0, 1.0, "fp32")["STOI"] - 1.0) < 1e-3
    assert evaluate(x, x, 16000, 1.0, 1.0, "fp32")["Bitrate (kbps)"] == round(BITRATE_KBPS, 4)
    assert evaluate(x, x, 16000, 1.0, 1.0, "lscodec-25hz", 0.25)["Bitrate (kbps)"] == 0.25
    print("core self-check ok")