"""The MP3 as a measurement: quantised MDCT coefficients, their quantiser cells, and the decoder's
synthesis as an exact linear operator (and its adjoint, which is the encoder-side analysis).

    bs  = read_bitstream("mix.mp3")            # runs ./mp3dump (minimp3 with a hook)
    syn = Synthesis(bs)                        # S: coded coefficients -> PCM, per channel
    pcm = syn.decode()                         # equals minimp3's own output
    lo, hi = cells(bs)                         # quantiser cell of every coded coefficient

The pipeline per channel (ISO 11172-3 / 13818-3 layer III):
    short-block reorder -> anti-alias butterflies -> IMDCT/window/overlap-add (per block-type pair)
    -> sign flip of odd samples in odd subbands -> 32-band polyphase synthesis.
The IMDCT/window/overlap and polyphase stages are captured as impulse responses of minimp3 itself,
so the operator matches the decoder to float precision. Mid/side: L = c0 + c1, R = c0 - c1
(minimp3 folds the 1/sqrt(2) into the step size).
"""
from __future__ import annotations

import os
import subprocess
import tempfile

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
DUMP = os.path.join(HERE, "mp3dump")
REC = 10 * 4 + 40 + 576 * 4 * 2
AA_C = np.array([-0.6, -0.535, -0.33, -0.185, -0.095, -0.041, -0.0142, -0.0037])
AA_CS = 1 / np.sqrt(1 + AA_C ** 2)
AA_CA = np.abs(AA_C) / np.sqrt(1 + AA_C ** 2)      # minimp3 writes the sign into the butterfly


def _build():
    if not os.path.exists(DUMP) or os.path.getmtime(DUMP) < os.path.getmtime(os.path.join(HERE, "mp3dump.c")):
        subprocess.run(["gcc", "-O2", "-o", DUMP, os.path.join(HERE, "mp3dump.c"), "-lm"], check=True)


def read_bitstream(path):
    _build()
    with tempfile.TemporaryDirectory() as td:
        gr, pcm = os.path.join(td, "x.gr"), os.path.join(td, "x.pcm")
        subprocess.run([DUMP, path, gr, pcm], check=True, capture_output=True)
        raw = np.fromfile(gr, dtype=np.uint8)
        pcm = np.fromfile(pcm, dtype=np.float32)
    n = raw.size // REC
    raw = raw[: n * REC].reshape(n, REC)
    head = raw[:, :40].copy().view(np.int32)
    assert np.all(head[:, 0] == 0x4D503347), "bad dump"
    sfb = raw[:, 40:80]
    x = raw[:, 80:80 + 2304].copy().view(np.float32)
    step = raw[:, 80 + 2304:].copy().view(np.float32)
    nch = int(head[:, 2].max()) + 1
    G = n // nch
    shp = lambda a: a.reshape(G, nch, *a.shape[1:])
    bs = {"hz": int(head[0, 7]), "nch": nch, "G": G,
          "block": shp(head[:, 3]), "mixed": shp(head[:, 4]), "ms": shp(head[:, 5])[:, 0], "is": shp(head[:, 6])[:, 0],
          "sfb": shp(sfb), "x": shp(x).astype(np.float64), "step": shp(step).astype(np.float64),
          "pcm": pcm.reshape(-1, nch).T.astype(np.float64)}
    assert not bs["mixed"].any(), "mixed blocks not supported"
    # integer quantised magnitude: |x| = |q|^(4/3) * step
    with np.errstate(divide="ignore", invalid="ignore"):
        q = np.where(bs["step"] > 0, np.round((np.abs(bs["x"]) / bs["step"]) ** 0.75), 0)
    bs["q"] = (np.sign(bs["x"]) * q).astype(np.int32)
    return bs


def _ir(mode):
    _build()
    with tempfile.TemporaryDirectory() as td:
        p = os.path.join(td, "ir")
        subprocess.run([DUMP, mode, p], check=True)
        return np.fromfile(p, dtype=np.float32).astype(np.float64)


def short_perm(widths):
    """Index map for the decoder's short-block reorder: out[j] = in[perm[j]] (first sum(3*w) entries)."""
    perm, src = [], 0
    w = [int(v) for v in widths if v]
    for b in range(0, len(w), 3):
        L = w[b]
        for i in range(L):
            perm += [src + i, src + L + i, src + 2 * L + i]
        src += 3 * L
    perm = np.array(perm, dtype=np.int64)
    full = np.arange(576)
    full[: len(perm)] = perm
    return full


class Synthesis:
    """S: coded coefficients (G, nch, 576) -> PCM (nch, 576*G), exactly as minimp3 decodes; S.T is its adjoint."""

    def __init__(self, bs):
        self.bs = bs
        G, nch = bs["G"], bs["nch"]
        ir = _ir("--imdct-ir").reshape(4, 4, 2, 18, 2, 18)          # prev, cur, parity, k, (same|next), n
        self.A = ir[:, 0, :, :, 0, :]                                # prev, parity, k, n   (same granule)
        self.B = ir[:, :, :, :, 1, :]                                # prev, cur, parity, k, n (next granule)
        h = _ir("--synth-ir").reshape(32, -1)                        # subband, samples
        L = int(np.max(np.nonzero(np.abs(h).max(axis=0) > 0)[0])) + 1
        h = h[:, :L]
        self.M = int(np.ceil(L / 32))
        h = np.pad(h, ((0, 0), (0, self.M * 32 - L)))
        self.H = h.reshape(32, self.M, 32).transpose(1, 2, 0)        # m, j(out sample), k(subband)
        self.perm = np.zeros((G, nch, 576), dtype=np.int64)
        for g in range(G):
            for c in range(nch):
                self.perm[g, c] = short_perm(bs["sfb"][g, c]) if bs["block"][g, c] == 2 else np.arange(576)
        self.inv = np.argsort(self.perm, axis=-1)

    # ── per-channel stages (X: (G, 576)) ──
    def _reorder(self, X, c):
        return np.take_along_axis(X, self.perm[:, c], axis=1)

    def _reorder_T(self, Y, c):
        return np.take_along_axis(Y, self.inv[:, c], axis=1)

    def _alias(self, X, c, transpose=False):
        X = X.copy().reshape(-1, 32, 18)
        long = self.bs["block"][:, c] != 2
        i = np.arange(8)
        Xl = X[long]
        lo, hi = Xl[:, :-1][:, :, 17 - i], Xl[:, 1:][:, :, i]          # d = lower subband tail, u = upper head
        if not transpose:   # u' = u cs - d ca ; d' = u ca + d cs
            nu, nd = hi * AA_CS - lo * AA_CA, hi * AA_CA + lo * AA_CS
        else:
            nu, nd = hi * AA_CS + lo * AA_CA, -hi * AA_CA + lo * AA_CS
        Xl[:, 1:, :8] = nu
        Xl[:, :-1, 17 - i] = nd
        X[long] = Xl
        return X.reshape(-1, 576)

    def _imdct(self, X, c):
        X = X.reshape(-1, 32, 18)
        bt = self.bs["block"][:, c]
        par = np.arange(32) % 2
        Y = np.einsum("gskn,gsk->gsn", self.A[bt][:, par], X)
        Bm = self.B[bt[:-1], bt[1:]][:, par]                          # (G-1, 32, 18, 18)
        Y[1:] += np.einsum("gskn,gsk->gsn", Bm, X[:-1])
        return Y

    def _imdct_T(self, Y, c):
        bt = self.bs["block"][:, c]
        par = np.arange(32) % 2
        X = np.einsum("gskn,gsn->gsk", self.A[bt][:, par], Y)
        Bm = self.B[bt[:-1], bt[1:]][:, par]
        X[:-1] += np.einsum("gskn,gsn->gsk", Bm, Y[1:])
        return X.reshape(-1, 576)

    def _poly(self, Z):
        S = Z.transpose(0, 2, 1).reshape(-1, 32)                     # slots, subband
        T = S.shape[0]
        out = np.zeros((T + self.M, 32))
        for m in range(self.M):
            out[m:m + T] += S @ self.H[m].T
        return out[:T].reshape(-1)

    def _poly_T(self, y):
        T = y.size // 32
        Y = np.zeros((T + self.M, 32))
        Y[:T] = y.reshape(T, 32)
        S = np.zeros((T, 32))
        for m in range(self.M):
            S += Y[m:m + T] @ self.H[m]
        return S.reshape(-1, 18, 32).transpose(0, 2, 1)

    def channel(self, X, c):
        return self._poly(self._imdct(self._alias(self._reorder(X, c), c), c))

    def channel_T(self, y, c):
        return self._reorder_T(self._alias(self._imdct_T(self._poly_T(y), c), c, transpose=True), c)

    # ── coded domain (mid/side) <-> left/right ──
    def to_lr(self, Xc):
        ms = self.bs["ms"].astype(bool)
        L, R = Xc[:, 0].copy(), Xc[:, 1].copy()
        L[ms], R[ms] = Xc[ms, 0] + Xc[ms, 1], Xc[ms, 0] - Xc[ms, 1]
        return L, R

    def __call__(self, Xc):
        """Coded coefficients (G, 2, 576) -> stereo PCM (2, 576 G)."""
        L, R = self.to_lr(Xc)
        return np.stack([self.channel(L, 0), self.channel(R, 1)])

    def T(self, y):
        """Adjoint: stereo PCM (2, 576 G) -> coded-domain coefficients (G, 2, 576)."""
        aL, aR = self.channel_T(y[0], 0), self.channel_T(y[1], 1)
        ms = self.bs["ms"].astype(bool)
        out = np.stack([aL, aR], axis=1)
        out[ms, 0], out[ms, 1] = aL[ms] + aR[ms], aL[ms] - aR[ms]
        return out

    def decode(self):
        return self(self.bs["x"])

    def gram_diag(self):
        """diag(S^T S) per coded coefficient (G, 2, 576). S is orthogonal up to these per-coefficient
        scales (long vs short IMDCT gain, x2 for mid/side), so the encoder's analysis is
        A = W S^T with W = 1 / gram_diag (exact away from the first and last granule)."""
        if hasattr(self, "_gd"):
            return self._gd
        G = self.bs["G"]
        par = np.arange(32) % 2
        poly = 128.0
        out = np.zeros((G, 2, 576))
        for c in range(2):
            bt = self.bs["block"][:, c]
            nxt = np.r_[bt[1:], bt[-1:]]
            nA = (self.A[bt][:, par] ** 2).sum(-1)                        # G, 32, 18 (post-reorder layout)
            nB = (self.B[bt, nxt][:, par] ** 2).sum(-1)
            d = (nA + nB).reshape(G, 576)
            out[:, c] = np.take_along_axis(d, self.inv[:, c], axis=1) * poly   # back to bitstream order
        ms = self.bs["ms"].astype(bool)
        out[ms] *= 2.0
        self._gd = out
        return out

    def analysis(self, y):
        """Encoder-side analysis A = W S^T: stereo PCM (2, 576 G) -> coded coefficients (G, 2, 576)."""
        return self.T(y) / self.gram_diag()


def cells(bs, slack=0.0):
    """Quantiser cell [lo, hi] of every coded coefficient, in the decoded scale.

    LAME rounds to the nearest dequantised level: q is chosen so |c|/step is closest to |q|^(4/3).
    So |c| lies between the midpoints of neighbouring levels:
        [ ((|q|-1)^(4/3) + |q|^(4/3)) / 2,  (|q|^(4/3) + (|q|+1)^(4/3)) / 2 ] * step,  with the sign of q;
    q = 0 gives |c| < step / 2. (At q = 0/1 this is ISO's 0.4054 rounding offset, 1 - 0.5^(3/4).)
    `slack` widens every cell by that fraction of a step on each side."""
    q, st = np.abs(bs["q"]).astype(np.float64), bs["step"]
    lv = lambda n: np.maximum(n, 0) ** (4 / 3)
    a = np.where(q > 0, (lv(q - 1) + lv(q)) / 2, 0.0) * st
    b = (lv(q) + lv(q + 1)) / 2 * st
    a, b = np.maximum(a - slack * st, 0), b + slack * st
    s = np.sign(bs["q"])
    lo = np.where(s > 0, a, np.where(s < 0, -b, -b))
    hi = np.where(s > 0, b, np.where(s < 0, -a, b))
    return lo, hi


def reliable(bs, lowpass_margin=18):
    """Which cells to trust: True where the coefficient was genuinely quantised.

    Two places where a coded zero says nothing about the signal: above the encoder's lowpass
    (highest bin that is ever nonzero, minus a subband of margin), and scalefactor bands that the
    encoder zeroed as a whole (it may drop a band instead of quantising it)."""
    q = bs["q"]
    nz = (q != 0).mean(axis=(0, 1))
    top = int(np.max(np.nonzero(nz > 0.01)[0])) if (nz > 0.01).any() else 575
    ok = (bs["step"] > 0) & (np.arange(576) < top - lowpass_margin)[None, None, :]
    G, C = q.shape[:2]
    for g in range(G):
        for c in range(C):
            start = 0
            for w in bs["sfb"][g, c].astype(int):
                if not w:
                    break
                if not q[g, c, start:start + w].any():
                    ok[g, c, start:start + w] = False
                start += w
    return ok