"""FocalCodec 12.5 Hz weight-quantization lab (Gradio / ZeroGPU)."""

import hashlib
import time

import gradio as gr
import numpy as np
import soundfile as sf
import spaces  # patches torch.cuda for ZeroGPU; must stay above `import torch`
import torch
from scipy.signal import resample_poly

import core
import int8nn
from core import QUANT_LEVELS

MODEL_REPO = "lucadellalib/focalcodec_12_5hz"
HUB_REPO = "lucadellalib/focalcodec"
LS_REPO = "cantabile-kwok/lscodec_25hz"
LS_PREFIX = "lscodec-"
MAX_SECONDS = 20
HEADERS = ["Level", "SI-SDR (dB)", "STOI", "LSD (dB)", "Model size (MB)", "Bitrate (kbps)", "RTF"]
DATATYPES = ["str", "number", "number", "number", "number", "number", "number"]
EXAMPLES = [
    ["assets/251-118436-0003.wav"],
    ["assets/p226_006.wav"],
    ["assets/5_1812_20170628135856.wav"],
]
LS_BITRATE_KBPS = 25 * 10 / 1000  # 25 tokens/s x 10-bit codebook

codec = torch.hub.load(HUB_REPO, "focalcodec", config=MODEL_REPO, trust_repo=True)
codec.eval().requires_grad_(False)
codec.to("cuda")
BASE_STATE = {k: v.detach().cpu().clone() for k, v in codec.state_dict().items()}
SIZES = {level: core.model_size_mb(BASE_STATE, level) for level in QUANT_LEVELS}
NATIVE_LEVEL = "int8-native"
try:
    # fresh instance, not deepcopy: a second copy of the live (CUDA) codec is not deep-copyable there
    codec_i8 = int8nn.apply_int8(
        torch.hub.load(HUB_REPO, "focalcodec", config=MODEL_REPO, trust_repo=True).eval().requires_grad_(False).cuda(),
        BASE_STATE,
    )
except Exception as exc:  # noqa: BLE001  a runtime quirk must not take the demo down
    codec_i8 = None
    print(f"[int8] native FocalCodec disabled: {exc!r}")
NATIVE_LEVELS = (NATIVE_LEVEL,) if codec_i8 is not None else ()
if codec_i8 is not None:
    SIZES[NATIVE_LEVEL] = core.model_size_mb(BASE_STATE, "int8")  # + fp32 scales, <1% overhead
FOCAL_LEVELS = QUANT_LEVELS + NATIVE_LEVELS  # focal rows; LS natives stay out (LS_PREFIX)
SAMPLE_RATE = codec.sample_rate_input
assert codec.sample_rate_output == SAMPLE_RATE, "demo assumes a 16/16 kHz codec"


def _boot_lscodec() -> dict | None:
    """LSCodec-25 Hz baseline: WavLM-Large (first 6 layers) prompt encoder + encoder/vocoder."""
    from scipy import signal

    if not hasattr(signal, "kaiser"):  # their pqmf.py imports scipy<1.15's top-level kaiser
        from scipy.signal.windows import kaiser

        signal.kaiser = kaiser

    from pathlib import Path

    import yaml
    from huggingface_hub import hf_hub_download
    from transformers import WavLMModel

    from lscodec.utils import load_model, load_vocoder

    files = {
        name: hf_hub_download(LS_REPO, name, repo_type="model")
        for name in (
            "encoder_config.yml",
            "vocoder_config.yml",
            "codebook.npy",
            "lscodec_encoder.pt",
            "lscodec_vocoder.pt",
        )
    }
    encoder_config = yaml.safe_load(Path(files["encoder_config.yml"]).read_text())
    vocoder_config = yaml.safe_load(Path(files["vocoder_config.yml"]).read_text())
    encoder_config["pretrain_codebook"] = files["codebook.npy"]
    vocoder_config["vq_codebook"] = files["codebook.npy"]

    enc = load_model(encoder_config, files["lscodec_encoder.pt"]).eval().requires_grad_(False)
    dec = load_vocoder(vocoder_config, files["lscodec_vocoder.pt"]).eval().requires_grad_(False)
    codebook = torch.as_tensor(np.load(files["codebook.npy"], allow_pickle=True))
    books = torch.nn.ModuleList(
        [torch.nn.Embedding.from_pretrained(b, freeze=True) for b in (codebook if codebook.ndim == 3 else [codebook])]
    )
    wavlm = WavLMModel.from_pretrained("microsoft/wavlm-large").eval().requires_grad_(False)
    wavlm.encoder.layers = wavlm.encoder.layers[:6]  # prompt features come from layers 0..5
    for model in (enc, dec, books, wavlm):
        model.cuda()  # ZeroGPU wants module-level cuda placement (docs: model loading)
    states = {
        key: {k: v.detach().cpu().clone() for k, v in mod.state_dict().items()}
        for key, mod in (("enc", enc), ("dec", dec), ("books", books))
    }
    out = {
        "enc": enc,
        "dec": dec,
        "books": books,
        "wavlm": wavlm,
        "enc_state": states["enc"],
        "dec_state": states["dec"],
        "books_state": states["books"],
        "repeat": bool(vocoder_config.get("repeat_input_tokens")),
        "out_sr": int(vocoder_config["sampling_rate"]),
    }
    for key, mod in (("enc", enc), ("dec", dec), ("books", books)):
        out[key + "_fp"] = mod  # fp32 module for state swaps; `_load_lscodec` restores it
        out[key + "_i8"] = None
    try:
        # fresh instances, not deepcopy: legacy weight-normed modules hold non-leaf tensors
        fresh_enc = load_model(encoder_config, files["lscodec_encoder.pt"]).eval().requires_grad_(False)
        fresh_dec = load_vocoder(vocoder_config, files["lscodec_vocoder.pt"]).eval().requires_grad_(False)
        fresh_books = torch.nn.ModuleList(
            [torch.nn.Embedding.from_pretrained(b, freeze=True) for b in (codebook if codebook.ndim == 3 else [codebook])]
        )
        out["enc_i8"] = int8nn.apply_int8(fresh_enc.cuda(), states["enc"])
        out["dec_i8"] = int8nn.apply_int8(fresh_dec.cuda(), states["dec"])
        out["books_i8"] = int8nn.apply_int8(fresh_books.cuda(), states["books"])
    except Exception as exc:  # noqa: BLE001  native LSCodec is optional, all-or-nothing
        out["enc_i8"] = out["dec_i8"] = out["books_i8"] = None
        print(f"[int8] native LSCodec disabled: {exc!r}")
    return out


try:
    LS = _boot_lscodec()
except Exception as exc:  # noqa: BLE001  a third-party dep breaking must not take the demo down
    LS = None
    print(f"[lscodec] baseline disabled: {exc!r}")

LS_LEVELS = tuple(LS_PREFIX + level for level in QUANT_LEVELS) if LS is not None else ()
LS_NATIVE_LEVELS = (
    (LS_PREFIX + NATIVE_LEVEL,)
    if LS is not None and all(LS[key + "_i8"] is not None for key in ("enc", "dec", "books"))
    else ()
)
BENCH_LEVELS = QUANT_LEVELS + LS_LEVELS + NATIVE_LEVELS + LS_NATIVE_LEVELS


def _lscodec_prompt(audio: torch.Tensor) -> torch.Tensor:
    """WavLM-Large layers 0..5 features: a function of the input only, so compute it once per clip."""
    with torch.inference_mode():
        prompt_input = torch.nn.functional.layer_norm(audio, audio.shape)  # fairseq WavLM cfg.normalize
        return LS["wavlm"](prompt_input, output_hidden_states=True).hidden_states[6]


def _lscodec_decode(prompt: torch.Tensor, audio: torch.Tensor, sr: int) -> np.ndarray:
    """Reconstruct `audio` with LSCodec at the currently loaded weight level, resampled back to `sr`."""
    with torch.inference_mode():
        _, _, index = LS["enc"].encode(audio.unsqueeze(0))
        if LS["repeat"]:
            index = index.repeat_interleave(2, dim=0)
        vqvec = torch.cat([book(index[:, i]) for i, book in enumerate(LS["books"])], dim=-1).unsqueeze(0)
        rec = LS["dec"].inference(vqvec, prompt)[-1].view(-1).float().cpu().numpy()
    out_sr = LS["out_sr"]
    if out_sr != sr:
        g = np.gcd(out_sr, sr)
        rec = resample_poly(rec, sr // g, out_sr // g).astype(np.float32)
    return rec


def _load_lscodec(level: str) -> None:
    """Point the LSCodec encoder, vocoder and codebook at `level` (QUANT_LEVELS entry or int8-native)."""
    if level == NATIVE_LEVEL:
        for key in ("enc", "dec", "books"):
            LS[key] = LS[key + "_i8"]
        return
    for key in ("enc", "dec", "books"):
        if LS[key] is not LS[key + "_fp"]:  # a previous native swap is active -> restore fp32 module
            LS[key] = LS[key + "_fp"]
        state = LS[key + "_state"]
        LS[key].load_state_dict(state if level == "fp32" else core.quantize_state(state, level))


def _lscodec_size_mb(level: str) -> float:
    """Stored size of the codec at `level`; the WavLM prompt encoder is excluded (see the tab note)."""
    grid = "int8" if level == NATIVE_LEVEL else level  # native twins share int8's storage grid
    return sum(core.model_size_mb(LS[key + "_state"], grid) for key in ("enc", "dec", "books"))


def _estimate_duration(audio_path: str, levels: tuple) -> float:
    try:
        seconds = min(sf.info(audio_path).duration, MAX_SECONDS)
    except Exception:  # noqa: BLE001  unreadable input falls back to the full cap
        seconds = float(MAX_SECONDS)
    native_slack = 2.5 * sum(1 for level in levels if level.endswith("native"))  # W8A8 rows run heavier
    # ZeroGPU rejects a request above the visitor's tier cap, so never ask for more than 60 s.
    return min(60.0, 3.0 + len(levels) + 1.5 * seconds + native_slack)


@spaces.GPU(duration=_estimate_duration)
def _encode_decode(audio_path: str, levels: tuple) -> dict:
    """Reconstruct `audio_path` once per weight-quantization level, scoring each result."""
    ref, sr = core.load_audio(audio_path, SAMPLE_RATE, MAX_SECONDS)
    seconds = len(ref) / sr
    rows, reconstructed = [], {}
    focal_levels = tuple(level for level in levels if level in FOCAL_LEVELS)
    ls_levels = tuple(level for level in levels if level.startswith(LS_PREFIX))
    if focal_levels:
        sig = torch.from_numpy(ref).unsqueeze(0).cuda()
        with torch.inference_mode():  # untimed warmup so the first row's RTF isn't inflated
            codec.sig_to_sig(sig)
            if NATIVE_LEVEL in focal_levels:
                codec_i8.sig_to_sig(sig)
        for level in focal_levels:
            model = codec_i8 if level == NATIVE_LEVEL else codec
            if level != NATIVE_LEVEL:
                codec.load_state_dict(BASE_STATE if level == "fp32" else core.quantize_state(BASE_STATE, level))
            torch.cuda.synchronize()
            start = time.perf_counter()
            with torch.inference_mode():
                rec = model.sig_to_sig(sig).squeeze(0).float().cpu().numpy()
            torch.cuda.synchronize()
            rtf = (time.perf_counter() - start) / seconds
            rows.append(core.evaluate(ref, rec, sr, rtf, SIZES[level], level))
            reconstructed[level] = (sr, rec)
        codec.load_state_dict(BASE_STATE)
    if ls_levels:
        audio = torch.from_numpy(ref).float().view(1, -1).cuda()
        with torch.inference_mode():  # one prompt for every level, plus an untimed warmup pass
            prompt = _lscodec_prompt(audio)
            _lscodec_decode(prompt, audio, sr)
            if any(level[len(LS_PREFIX) :] == NATIVE_LEVEL for level in ls_levels):
                _load_lscodec(NATIVE_LEVEL)
                _lscodec_decode(prompt, audio, sr)
                _load_lscodec("fp32")
        for level in ls_levels:
            _load_lscodec(level[len(LS_PREFIX) :])
            torch.cuda.synchronize()
            start = time.perf_counter()
            rec = _lscodec_decode(prompt, audio, sr)
            torch.cuda.synchronize()
            rtf = (time.perf_counter() - start) / seconds
            base = level[len(LS_PREFIX) :]
            rows.append(core.evaluate(ref, rec, sr, rtf, _lscodec_size_mb(base), level, LS_BITRATE_KBPS))
            reconstructed[level] = (sr, rec)
        _load_lscodec("fp32")  # leave the baseline weights in place for the next request
    return {"input": (sr, ref), "rows": rows, "audio": reconstructed}


def _table(rows: list[dict]) -> list[list]:
    return [[row[header] for header in HEADERS] for row in rows]


_CACHE: dict[str, dict] = {}


def run(audio_path: str, levels: list | tuple) -> dict:
    """Encode and decode the audio at every requested level (either codec), cached per level set."""
    if not audio_path:
        raise gr.Error("Provide audio first: upload, record, or pick an example.")
    levels = tuple(levels)
    with open(audio_path, "rb") as fh:
        key = hashlib.md5(fh.read()).hexdigest() + "|" + ",".join(levels)
    if key not in _CACHE:
        if len(_CACHE) >= 6:
            _CACHE.clear()
        _CACHE[key] = _encode_decode(audio_path, levels)
    return _CACHE[key]


def compare(audio_path: str):
    """Rebuild the audio at all eight levels and return one player + spectrogram per level."""
    result = run(audio_path, QUANT_LEVELS)
    audio = [result["audio"][level] for level in QUANT_LEVELS]
    plots = [
        core.spectrogram_figure(result["input"][1], result["audio"][level][1], SAMPLE_RATE, level)
        for level in QUANT_LEVELS
    ]
    return audio + plots


def playground(audio_path: str, level: str):
    """Rebuild the audio at a single quantization level, with its metrics row."""
    result = run(audio_path, [level])
    rec = result["audio"][level]
    return (
        rec,
        core.spectrogram_figure(result["input"][1], rec[1], SAMPLE_RATE, level),
        _table(result["rows"]),
    )


def benchmark(audio_path: str):
    """Score every level of both codecs on one input: metrics table plus summary chart."""
    rows = run(audio_path, BENCH_LEVELS)["rows"]
    return _table(rows), core.metrics_figure(rows)


HEADER = f"""# ⚡ FocalCodec 12.5 Hz — weight quantization lab

Same checkpoint ([`lucadellalib/focalcodec_12_5hz`](https://huggingface.co/lucadellalib/focalcodec_12_5hz)),
eight numeric formats for its **{sum(v.numel() for v in BASE_STATE.values()) / 1e6:.0f} M weights**: `int4 … int32` and `fp4 … fp32`.
The bitrate never changes — **{core.BITRATE_KBPS:g} kbps** (12.5 tokens/s × 13 bits) — only the storage precision of the network itself.

**Method:** post-training fake quantization, per-output-channel absmax scaling, round-to-nearest, dequantized back to fp32 for inference. `fp4` = E2M1, `fp8` = E4M3, `fp16`/`fp32` are exact casts. No fine-tuning.

**`int8-native` rows** run the same int8 weights through a real **W8A8 runtime** instead: weights stay int8 in memory (never dequantized), each GEMM quantizes its activations per row on the fly and computes `int8 × int8 → int32` (`torch._int_mm`).

⚠️ *SI-SDR sits far below 0 dB at every level: at 0.16 kbps the codec does not preserve waveform phase, and the FocalCodec paper reports SI-SNR ≈ −37 dB for the untouched model. It is an anchor, not a quality score — read **STOI** (intelligibility) and **LSD** (log-spectral distance, lower is better) instead.*
"""

FOOTER = """Audio examples from the [FocalCodec repository](https://github.com/lucadellalib/focalcodec)
(LibriSpeech, VoiceBank, MagicData). Model and method by Luca Della Libera et al., Apache-2.0.
The **`lscodec-…`** rows are [LSCodec](https://github.com/X-LANCE/LSCodec-Inference) by
Yiwei Guo et al. (MIT), quantized identically and self-prompted with WavLM-Large from
[`microsoft/wavlm-large`](https://huggingface.co/microsoft/wavlm-large)."""

with gr.Blocks(title="FocalCodec quantization lab") as demo:
    gr.Markdown(HEADER)

    with gr.Tab("Compare"):
        with gr.Row():
            with gr.Column(scale=1, min_width=320):
                cmp_input = gr.Audio(
                    sources=["upload", "microphone"], type="filepath", label="Input audio"
                )
                gr.Examples(examples=EXAMPLES, inputs=cmp_input, label="Examples")
                cmp_button = gr.Button("Reconstruct at all 8 levels", variant="primary")
                gr.Markdown(f"Input capped at {MAX_SECONDS} s.")
            with gr.Column(scale=4):
                cmp_audio, cmp_plot = {}, {}
                for chunk in (QUANT_LEVELS[:4], QUANT_LEVELS[4:]):
                    with gr.Row():
                        for level in chunk:
                            with gr.Column():
                                gr.Markdown(f"**{level}**")
                                cmp_audio[level] = gr.Audio(interactive=False, label=level)
                                cmp_plot[level] = gr.Plot(label=f"{level} spectrogram")
        cmp_button.click(
            compare,
            inputs=cmp_input,
            outputs=[cmp_audio[level] for level in QUANT_LEVELS]
            + [cmp_plot[level] for level in QUANT_LEVELS],
        )

    with gr.Tab("Playground"):
        with gr.Row():
            with gr.Column(scale=1, min_width=320):
                pg_input = gr.Audio(
                    sources=["upload", "microphone"], type="filepath", label="Input audio"
                )
                gr.Examples(examples=EXAMPLES, inputs=pg_input, label="Examples")
                pg_level = gr.Dropdown(
                    choices=list(BENCH_LEVELS),
                    value="int8",
                    label="Codec / weight format",
                )
                pg_button = gr.Button("Encode + decode", variant="primary")
            with gr.Column(scale=3):
                pg_output = gr.Audio(interactive=False, label="Reconstruction")
                pg_plot = gr.Plot(label="Input vs reconstruction")
        pg_table = gr.Dataframe(headers=HEADERS, datatype=DATATYPES, interactive=False)
        pg_button.click(
            playground, inputs=[pg_input, pg_level], outputs=[pg_output, pg_plot, pg_table]
        )

    with gr.Tab("Benchmark"):
        with gr.Row():
            with gr.Column(scale=1, min_width=320):
                bm_input = gr.Audio(
                    sources=["upload", "microphone"], type="filepath", label="Input audio"
                )
                gr.Examples(examples=EXAMPLES, inputs=bm_input, label="Examples")
                bm_button = gr.Button(f"Run all {len(BENCH_LEVELS)} levels", variant="primary")
            with gr.Column(scale=3):
                gr.Markdown(
                    "**STOI** ↑ intelligibility · **LSD** ↓ spectral error · "
                    "**RTF** = encode+decode time ÷ audio duration · size assumes "
                    "the stored weight format."
                )
                if LS_LEVELS:
                    gr.Markdown(
                        f"Rows prefixed `lscodec-` are the second codec, quantized the same way: "
                        f"it runs at **{LS_BITRATE_KBPS:g} kbps** (25 tokens/s × 10 bits) versus "
                        f"FocalCodec's **{core.BITRATE_KBPS:g} kbps**, and its size covers the encoder, "
                        "vocoder and codebook only — the WavLM-Large prompt encoder it needs is 1.2 GB "
                        "and stays excluded."
                    )
                if NATIVE_LEVELS:
                    gr.Markdown(
                        f"`{NATIVE_LEVEL}` and `{LS_PREFIX}{NATIVE_LEVEL}` are the same int8 weights run "
                        "natively (W8A8-dynamic, int8-resident) — their STOI/LSD should track the plain "
                        "`int8` rows, but the weights are never dequantized."
                    )
        bm_table = gr.Dataframe(headers=HEADERS, datatype=DATATYPES, interactive=False)
        bm_plot = gr.Plot(label="Quality and size per level")
        bm_button.click(benchmark, inputs=bm_input, outputs=[bm_table, bm_plot])

    gr.Markdown(FOOTER)

demo.launch(mcp_server=True, theme=gr.themes.Soft())