"""Post-training INT8 quantization for the PhonemeTCN, targeting a small
custom streaming engine on the ESP32-S3.

Scheme (classic TFLite-style, but explicit and portable):
  - weights: symmetric per-output-channel int8, BN folded into conv
  - activations: per-tensor asymmetric-free symmetric int8 (ReLU outputs
    use unsigned range via zero offset 0..127 semantics kept simple:
    symmetric [-127,127] everywhere)
  - accumulators: int32; requantization by fixed-point multiplier per layer

Steps:
  python -m phoneme_engine.quantize calibrate   # activation ranges
  python -m phoneme_engine.quantize verify      # int8 sim vs float parity
  python -m phoneme_engine.quantize export      # C header + weights bin
"""

import json
import sys
from pathlib import Path

import numpy as np
import torch

from .data import LibriPhonemes, collate
from .features import LogMel
from .model import PhonemeTCN

ROOT = Path(__file__).resolve().parent.parent
QDIR = ROOT / "export"


def fold_bn(conv_w, conv_b, bn):
    """Fold BatchNorm into conv weights/bias (eval-mode running stats)."""
    gamma = bn.weight.detach().numpy()
    beta = bn.bias.detach().numpy()
    mean = bn.running_mean.detach().numpy()
    var = bn.running_var.detach().numpy()
    scale = gamma / np.sqrt(var + bn.eps)
    w = conv_w * scale[:, None, None]
    b = (conv_b if conv_b is not None else 0.0) * scale + beta - mean * scale
    return w, b


def extract_layers(model):
    """Flatten the model into a list of layer dicts with folded BN.

    Layer kinds: conv (dense conv1d), dw (depthwise), each with
    weight (out, in, k) / (ch, 1, k), bias, stride, dilation, plus
    residual bookkeeping: blocks add their input to the pw output.
    """
    layers = []
    w, b = fold_bn(model.stem.conv.weight.detach().numpy(),
                   None if model.stem.conv.bias is None
                   else model.stem.conv.bias.detach().numpy(),
                   model.stem_bn)
    layers.append(dict(kind="conv", w=w, b=b, stride=2, dilation=1,
                       relu=True, residual=False))
    for blk in model.blocks:
        w, b = fold_bn(blk.dw.conv.weight.detach().numpy(),
                       None if blk.dw.conv.bias is None
                       else blk.dw.conv.bias.detach().numpy(),
                       blk.bn1)
        layers.append(dict(kind="dw", w=w, b=b, stride=1,
                           dilation=blk.dw.conv.dilation[0], relu=True,
                           residual=False))
        w, b = fold_bn(blk.pw.weight.detach().numpy(),
                       None if blk.pw.bias is None
                       else blk.pw.bias.detach().numpy(),
                       blk.bn2)
        # pw output adds the block input, then ReLU
        layers.append(dict(kind="conv", w=w, b=b, stride=1, dilation=1,
                           relu=True, residual=True))
    layers.append(dict(kind="conv",
                       w=model.head.weight.detach().numpy(),
                       b=model.head.bias.detach().numpy(),
                       stride=1, dilation=1, relu=False, residual=False))
    return layers


def float_forward(layers, feats):
    """Reference float forward pass on (T, 40) features using the flat
    layer list. Must match the PyTorch model exactly (verified)."""
    x = feats.T  # (C, T)
    block_input = None
    for lay in layers:
        if lay["kind"] == "dw":
            block_input = x  # residual adds the block's input, saved here
        w, b = lay["w"], lay["b"]
        k = w.shape[2]
        d = lay["dilation"]
        pad = d * (k - 1)
        xin = np.pad(x, ((0, 0), (pad, 0)))
        T = x.shape[1]
        out_T = (T - 1) // lay["stride"] + 1
        out_C = w.shape[0]
        y = np.zeros((out_C, out_T), dtype=np.float64)
        for t in range(out_T):
            base = t * lay["stride"] + pad
            taps = xin[:, [base - d * (k - 1 - i) for i in range(k)]]
            if lay["kind"] == "dw":
                y[:, t] = (taps * w[:, 0, :]).sum(axis=1) + b
            else:
                y[:, t] = np.tensordot(w, taps, axes=([1, 2], [0, 1])) + b
        if lay["residual"]:
            y = y + block_input[:, :out_T]
        if lay["relu"]:
            y = np.maximum(y, 0.0)
        x = y
    return x.T  # (T, classes)


def collect_calibration_feats(n_utts=64):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    frontend = LogMel().to(device).eval()
    ds = LibriPhonemes(str(ROOT / "data"), "dev-clean")
    feats = []
    with torch.no_grad():
        for i in range(0, n_utts * 40, 40):
            wav, _ = ds[i % len(ds)]
            f = frontend(wav.unsqueeze(0).to(device))[0].cpu().numpy()
            feats.append(f[:400])  # up to 4 s per utterance
            if len(feats) >= n_utts:
                break
    return feats


def calibrate_scales(layers, feats_list, pctl=99.95):
    """Per-layer activation scales from representative audio: runs the
    float replay and records robust max-abs at every layer boundary."""
    n_layers = len(layers)
    maxima = [[] for _ in range(n_layers + 1)]  # +1 for the input feats
    for feats in feats_list:
        maxima[0].append(np.percentile(np.abs(feats), pctl))
        x = feats.T
        block_input = None
        for li, lay in enumerate(layers):
            if lay["kind"] == "dw":
                block_input = x
            w, b = lay["w"], lay["b"]
            k = w.shape[2]
            d = lay["dilation"]
            pad = d * (k - 1)
            xin = np.pad(x, ((0, 0), (pad, 0)))
            out_T = (x.shape[1] - 1) // lay["stride"] + 1
            y = np.zeros((w.shape[0], out_T))
            for t in range(out_T):
                base = t * lay["stride"] + pad
                taps = xin[:, [base - d * (k - 1 - i) for i in range(k)]]
                if lay["kind"] == "dw":
                    y[:, t] = (taps * w[:, 0, :]).sum(axis=1) + b
                else:
                    y[:, t] = np.tensordot(w, taps,
                                           axes=([1, 2], [0, 1])) + b
            if lay["residual"]:
                y = y + block_input[:, :out_T]
            if lay["relu"]:
                y = np.maximum(y, 0.0)
            maxima[li + 1].append(np.percentile(np.abs(y), pctl))
            x = y
    scales = [max(float(np.max(m)), 1e-3) / 127.0 for m in maxima]
    return scales


def quantize_weights(layers):
    """Symmetric per-output-channel int8 weights; returns quantized copies
    (float values on the int8 grid) plus the raw int8 arrays and scales."""
    qlayers = []
    for lay in layers:
        w = lay["w"]
        s_w = np.abs(w).reshape(w.shape[0], -1).max(axis=1) / 127.0
        s_w = np.maximum(s_w, 1e-8)
        w_int = np.clip(np.round(w / s_w[:, None, None]), -127, 127)
        q = dict(lay)
        q["w"] = w_int * s_w[:, None, None]
        q["w_int"] = w_int.astype(np.int8)
        q["s_w"] = s_w
        qlayers.append(q)
    return qlayers


def fake_quant_forward(qlayers, scales, feats):
    """Float replay with activations snapped to the int8 grid at every
    layer boundary -- numerically equivalent to the integer engine."""
    def snap(x, s):
        return np.clip(np.round(x / s), -127, 127) * s

    x = snap(feats.T, scales[0])
    block_input = None
    block_input_scale = None
    for li, lay in enumerate(qlayers):
        if lay["kind"] == "dw":
            block_input = x
            block_input_scale = scales[li]
        w, b = lay["w"], lay["b"]
        k = w.shape[2]
        d = lay["dilation"]
        pad = d * (k - 1)
        xin = np.pad(x, ((0, 0), (pad, 0)))
        out_T = (x.shape[1] - 1) // lay["stride"] + 1
        y = np.zeros((w.shape[0], out_T))
        for t in range(out_T):
            base = t * lay["stride"] + pad
            taps = xin[:, [base - d * (k - 1 - i) for i in range(k)]]
            if lay["kind"] == "dw":
                y[:, t] = (taps * w[:, 0, :]).sum(axis=1) + b
            else:
                y[:, t] = np.tensordot(w, taps, axes=([1, 2], [0, 1])) + b
        if lay["residual"]:
            y = y + snap(block_input[:, :out_T], block_input_scale)
        if lay["relu"]:
            y = np.maximum(y, 0.0)
        x = snap(y, scales[li + 1])
    return x.T


def main():
    cmd = sys.argv[1] if len(sys.argv) > 1 else "verify"
    QDIR.mkdir(exist_ok=True)
    device = "cpu"
    model = PhonemeTCN().eval()
    state = torch.load(ROOT / "checkpoints" / "best.pt",
                       map_location=device, weights_only=True)
    model.load_state_dict(state["model"])
    layers = extract_layers(model)
    print(f"{len(layers)} layers extracted")

    if cmd == "sanity":
        feats = collect_calibration_feats(2)
        x = torch.from_numpy(feats[0]).float().unsqueeze(0)
        with torch.no_grad():
            ref = model(x)[0].numpy()
        ours = float_forward(layers, feats[0])
        err = np.abs(ref - ours).max()
        print(f"float reference vs pytorch max abs err: {err:.2e}")
        assert err < 1e-3, "layer extraction is wrong"
        print("sanity OK")

    elif cmd == "calibrate":
        feats_list = collect_calibration_feats(48)
        print(f"calibrating on {len(feats_list)} utterances...")
        scales = calibrate_scales(layers, feats_list)
        (QDIR / "act_scales.json").write_text(json.dumps(scales))
        print("activation scales:", [round(s, 5) for s in scales])
        print(f"saved -> {QDIR / 'act_scales.json'}")

    elif cmd == "verify":
        scales = json.loads((QDIR / "act_scales.json").read_text())
        qlayers = quantize_weights(layers)
        from .decoder import KeywordSpotter
        from .spot_file import load_wav
        from .features import LogMel
        device = "cuda" if torch.cuda.is_available() else "cpu"
        frontend = LogMel().to(device).eval()

        agree = tot = 0
        deltas = []
        clips = [(f, "sakura") for f in
                 sorted((ROOT / "test_clips" / "sakura_piper").glob("*.wav"))[:8]]
        clips += [(f, "hey orbit") for f in
                  sorted((ROOT / "test_clips" / "hey_orbit_piper").glob("*.wav"))[:8]]
        clips.append((ROOT / "diag_last.wav", "hey orbit"))
        for f, phrase in clips:
            wav = load_wav(str(f), device)
            with torch.no_grad():
                feats = frontend(wav)[0].cpu().numpy()
            lf = float_forward(layers, feats)
            lq = fake_quant_forward(qlayers, scales, feats)
            agree += (lf.argmax(1) == lq.argmax(1)).sum()
            tot += lf.shape[0]
            # spotting parity: decoder consumes logit differences directly
            sp = KeywordSpotter(phrase)
            sf_ = sp.best_score(lf - lf.max(axis=1, keepdims=True))
            sq_ = sp.best_score(lq - lq.max(axis=1, keepdims=True))
            if np.isfinite(sf_) or np.isfinite(sq_):
                deltas.append(abs(sf_ - sq_))
            print(f"  {f.name} [{phrase}]: float {sf_:7.2f}  int8 {sq_:7.2f}")
        print(f"frame argmax agreement: {agree/tot:.4f}")
        finite = [d for d in deltas if np.isfinite(d)]
        print(f"spot score delta (finite pairs): max {max(finite):.3f} "
              f"mean {np.mean(finite):.3f}; gate flips: "
              f"{len(deltas) - len(finite)}")

    elif cmd == "export":
        scales = json.loads((QDIR / "act_scales.json").read_text())
        qlayers = quantize_weights(layers)
        from .phones import PHONES
        lines = ["// Auto-generated by phoneme_engine.quantize export",
                 "// PhonemeTCN int8 weights + scales for the streaming",
                 "// wake word engine. Do not edit by hand.",
                 "#pragma once", "#include <stdint.h>", ""]
        lines.append(f"#define PWW_NUM_LAYERS {len(qlayers)}")
        lines.append(f"#define PWW_NUM_CLASSES {len(PHONES) + 1}")
        lines.append(f"#define PWW_INPUT_SCALE {scales[0]:.8f}f")
        lines.append(f"#define PWW_LOGIT_SCALE {scales[-1]:.8f}f")
        lines.append("")
        phones_str = ", ".join(f'"{p}"' for p in PHONES)
        lines.append(f"static const char *PWW_PHONES[] = {{{phones_str}}};")
        lines.append("")
        meta_rows = []
        total_bytes = 0
        for li, lay in enumerate(qlayers):
            w_int = lay["w_int"]
            out_c, in_c, k = w_int.shape
            flat = w_int.flatten()
            total_bytes += flat.size
            arr = ", ".join(str(int(v)) for v in flat)
            lines.append(f"static const int8_t PWW_W{li}[] = {{{arr}}};")
            # combined scale per channel: s_in * s_w[c]  (float requant)
            comb = scales[li] * lay["s_w"]
            arr = ", ".join(f"{v:.8e}f" for v in comb)
            lines.append(f"static const float PWW_S{li}[] = {{{arr}}};")
            arr = ", ".join(f"{v:.8e}f" for v in lay["b"])
            lines.append(f"static const float PWW_B{li}[] = {{{arr}}};")
            # depthwise weights are stored (ch, 1, k): the layer's true
            # input width is out_c, not the stored dim
            eff_in = out_c if lay["kind"] == "dw" else in_c
            meta_rows.append(
                f"  {{{1 if lay['kind'] == 'dw' else 0}, {eff_in}, {out_c}, "
                f"{k}, {lay['stride']}, {lay['dilation']}, "
                f"{1 if lay['relu'] else 0}, {1 if lay['residual'] else 0}, "
                f"PWW_W{li}, PWW_S{li}, PWW_B{li}, "
                f"{scales[li + 1]:.8f}f}}")
            lines.append("")
        lines.append(
            "typedef struct { uint8_t is_dw; uint16_t in_c, out_c; "
            "uint8_t k, stride, dilation, relu, residual; "
            "const int8_t *w; const float *s; const float *b; "
            "float out_scale; } pww_layer_t;")
        lines.append("")
        lines.append("static const pww_layer_t PWW_LAYERS[] = {")
        lines.append(",\n".join(meta_rows))
        lines.append("};")
        path = QDIR / "model_int8.h"
        path.write_text("\n".join(lines), encoding="utf-8")
        print(f"exported {total_bytes/1024:.0f} KB of int8 weights "
              f"-> {path}")

    elif cmd == "export-frontend":
        # exact DSP constants from the training frontend, so the C mel
        # frontend is identical by construction
        from .features import (HOP_LENGTH, N_FFT, N_MELS, SAMPLE_RATE,
                               WIN_LENGTH, LogMel)
        fe = LogMel()
        fb = fe.mel.mel_scale.fb.numpy()          # (n_freqs, n_mels)
        win = torch.hann_window(WIN_LENGTH, periodic=True).numpy()
        lines = ["// Auto-generated: mel frontend constants (exact copy of",
                 "// the training features). Do not edit.",
                 "#pragma once", ""]
        lines.append(f"#define PWW_FE_SR {SAMPLE_RATE}")
        lines.append(f"#define PWW_FE_NFFT {N_FFT}")
        lines.append(f"#define PWW_FE_WIN {WIN_LENGTH}")
        lines.append(f"#define PWW_FE_HOP {HOP_LENGTH}")
        lines.append(f"#define PWW_FE_NMELS {N_MELS}")
        lines.append(f"#define PWW_FE_NFREQS {fb.shape[0]}")
        lines.append("#define PWW_FE_EMA_ALPHA 0.02f")
        lines.append("")
        arr = ", ".join(f"{v:.8e}f" for v in win)
        lines.append(f"static const float PWW_FE_HANN[] = {{{arr}}};")
        arr = ", ".join(f"{v:.8e}f" for v in fb.T.flatten())
        lines.append("// mel filterbank, row-major (n_mels, n_freqs)")
        lines.append(f"static const float PWW_FE_MELFB[] = {{{arr}}};")
        path = QDIR / "frontend_data.h"
        path.write_text("\n".join(lines), encoding="utf-8")
        print(f"exported frontend constants -> {path}")


if __name__ == "__main__":
    main()