"""The integer super-resolution derivations of extensions/exact-upscale.md §3, written from
that text in integer arithmetic of their own, standard library only: the `WSR1` prior (§3.1),
`wai-upscale:sr2` (§3.2) and `wai-upscale:sr4` with its limits (§3.3). Shared by
`upscale_verify.py` and `synthesis_verify.py`; it shares no code with wai-rs.
"""
import hashlib
import struct

SR2, SR4 = "wai-upscale:sr2", "wai-upscale:sr4"
FACTOR = {SR2: 2, SR4: 4}
CAP = 1 << 26


class Unreadable(Exception):
    """A prior that breaks §3.1, or a picture past §3.3's limits: no output."""

    def __init__(self, code):
        super().__init__(code)
        self.code = code


def read_prior(b):
    """§3.1, in its order: `bad_magic`, `version`, `field_range`, `truncated`, `trailing`."""
    def need(at, n):
        if len(b) < at + n:
            raise Unreadable("truncated")
    need(0, 4)
    if b[:4] != b"WSR1":
        raise Unreadable("bad_magic")
    need(4, 1)
    if b[4] != 1:
        raise Unreadable("version")
    need(5, 3)
    n, act_hi = b[5], struct.unpack_from("<H", b, 6)[0]
    if not 1 <= n <= 8 or not 1 <= act_hi <= 32767:
        raise Unreadable("field_range")
    at, cin, layers = 8, 3, []
    for i in range(n):
        need(at, 4)
        k, shift, cout = b[at], b[at + 1], struct.unpack_from("<H", b, at + 2)[0]
        last = i + 1 == n
        if k not in (1, 3, 5) or shift > 31 or (cout != 12 if last else not 1 <= cout <= 64):
            raise Unreadable("field_range")
        at += 4
        nw = cout * cin * k * k
        need(at, nw + 6 * cout)
        w = list(struct.unpack_from("<%db" % nw, b, at))
        at += nw
        bias = list(struct.unpack_from("<%di" % cout, b, at))
        at += 4 * cout
        mult = list(struct.unpack_from("<%dH" % cout, b, at))
        at += 2 * cout
        if 0 in mult:
            raise Unreadable("field_range")
        layers.append({"k": k, "shift": shift, "cin": cin, "cout": cout, "w": w, "bias": bias, "mult": mult})
        cin = cout
    if at != len(b):
        raise Unreadable("trailing")
    return {"act_hi": act_hi, "layers": layers}


def up2(px, w, h):
    """§3.2 step 1: integer-operators §5E E11, the last row and column repeated."""
    out = bytearray(4 * w * h * 3)
    ow = 2 * w
    for y in range(h):
        for x in range(w):
            x1, y1 = min(x + 1, w - 1), min(y + 1, h - 1)
            for c in range(3):
                a, b2 = px[(y * w + x) * 3 + c], px[(y * w + x1) * 3 + c]
                c2, d = px[(y1 * w + x) * 3 + c], px[(y1 * w + x1) * 3 + c]
                out[((2 * y) * ow + 2 * x) * 3 + c] = a
                out[((2 * y) * ow + 2 * x + 1) * 3 + c] = (a + b2 + 1) >> 1
                out[((2 * y + 1) * ow + 2 * x) * 3 + c] = (a + c2 + 1) >> 1
                out[((2 * y + 1) * ow + 2 * x + 1) * 3 + c] = (a + b2 + c2 + d + 2) >> 2
    return bytes(out)


def half_up(p, s):
    """integer-operators §2 R_half_up: ⌊(p + 2^(s−1)) / 2^s⌋, p for s = 0."""
    return p if s == 0 else (p + (1 << (s - 1))) >> s


def sr2(prior, px, w, h):
    """§3.2 over a `w × h` RGB8 picture: `2w × 2h` RGB8. Refuses a picture past §3.3's limits."""
    if w < 1 or h < 1 or len(px) != w * h * 3:
        raise Unreadable("geometry")
    widest = max(l["cout"] for l in prior["layers"])
    if widest * w * h > CAP or 12 * w * h > CAP:
        raise Unreadable("too_large")
    n = w * h
    # Step 2: centred, channel-major.
    x = [[px[i * 3 + c] - 128 for i in range(n)] for c in range(3)]
    for li, l in enumerate(prior["layers"]):
        k, r, cin = l["k"], l["k"] // 2, l["cin"]
        last = li + 1 == len(prior["layers"])
        y = []
        for o in range(l["cout"]):
            taps = []
            for ci in range(cin):
                for ky in range(k):
                    for kx in range(k):
                        wt = l["w"][((o * cin + ci) * k + ky) * k + kx]
                        if wt:
                            taps.append((wt, ci, ky - r, kx - r))
            plane = []
            for yy in range(h):
                for xx in range(w):
                    a = 0
                    for wt, ci, dy, dx in taps:
                        sy, sx = yy + dy, xx + dx
                        if 0 <= sy < h and 0 <= sx < w:
                            a += wt * x[ci][sy * w + sx]
                    v = half_up((a + l["bias"][o]) * l["mult"][o], l["shift"])
                    plane.append(v if last else min(max(v, 0), prior["act_hi"]))
            y.append(plane)
        x = y
    base = up2(px, w, h)
    ow, oh = 2 * w, 2 * h
    out = bytearray(ow * oh * 3)
    for c in range(3):
        for oy in range(oh):
            for ox in range(ow):
                d = x[4 * c + 2 * (oy % 2) + ox % 2][(oy // 2) * w + ox // 2]
                i = (oy * ow + ox) * 3 + c
                out[i] = min(max(base[i] + d, 0), 255)
    return bytes(out)


_CACHE = {}


def derive(derivation, prior_bytes, px, w, h):
    """`wai-upscale:sr2` or `wai-upscale:sr4` (§3.3) under a prior's bytes: the output, or
    Unreadable. Results are kept by the prior's digest and the input, since a verifier meets
    one derivation many times."""
    key = (derivation, hashlib.sha256(prior_bytes).digest(), bytes(px), w, h)
    if key not in _CACHE:
        try:
            p = read_prior(prior_bytes)
            out = sr2(p, px, w, h)
            if derivation == SR4:
                out = sr2(p, out, 2 * w, 2 * h)
            elif derivation != SR2:
                raise Unreadable("derivation")
            _CACHE[key] = out
        except Unreadable as e:
            _CACHE[key] = e
    r = _CACHE[key]
    if isinstance(r, Unreadable):
        raise r
    return r