#!/usr/bin/env python3
"""Train the integer super-resolution priors of `extensions/exact-upscale.md` (§3), and
write the generated pictures the measurement reads.

    python tools/wai_int_sr_train.py --out <dir> [--width 16] [--layers 3] [--steps 4000]

Needs torch and numpy. It writes, under <dir>:

- `pictures/manifest.json` and `pictures/<name>.rgb`: the pictures, each RGB8 (rows top to
  bottom, R, G, B bytes), generated here from a seeded formula (`render`), never read from
  a file: a training set, a held-out set, and titles of eight frames each (a scene panned
  under a moving window). `--source <png>` adds one more title, a picture the caller
  supplies (the measurement adds the pinned crop of `tools/wai_sota_source_image.py`);
- `priors/generic.wsr`: a `WSR1` prior (§3.1) trained on the training set;
- `priors/<title>.wsr`: for each title, the generic prior fine-tuned on that title's frames:
  a per-title prior (§6);
- `report.json`: the seeds, the steps, the losses, and for each held-out picture and title
  the PSNR of the pinned upscale (§3.2 step 1) and of each prior's derivation, computed by the
  integer arithmetic of §3.2 (`int_forward`), not by the float model.

Training is in float (torch), then quantised after training and fine-tuned with the
quantisation in the loop. Every seed is fixed (`--seed`), but float training is not claimed
to be bit-reproducible: its sums depend on the BLAS and thread count of the machine that
runs it. What is pinned is each prior's bytes, by SHA-256. The derivation over those bytes is
exact (§3.2), whatever machine trained them.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import struct
from pathlib import Path

import numpy as np

# ---------------------------------------------------------------------------------------
# The pinned arithmetic of §3.2, in numpy int64: the reference this tool reports with.
# ---------------------------------------------------------------------------------------


def up2(lr: np.ndarray) -> np.ndarray:
    """`up2_predict` (integer-operators §5E E11): [h, w, 3] u8 -> [2h, 2w, 3] u8."""
    a = lr.astype(np.int64)
    b = np.concatenate([a[:, 1:], a[:, -1:]], axis=1)
    c = np.concatenate([a[1:], a[-1:]], axis=0)
    d = np.concatenate([b[1:], b[-1:]], axis=0)
    h, w, _ = a.shape
    out = np.empty((2 * h, 2 * w, 3), np.int64)
    out[0::2, 0::2] = a
    out[0::2, 1::2] = (a + b + 1) >> 1
    out[1::2, 0::2] = (a + c + 1) >> 1
    out[1::2, 1::2] = (a + b + c + d + 2) >> 2
    return out.astype(np.uint8)


def box_down(hr: np.ndarray) -> np.ndarray:
    """The low-resolution picture the measurement sends: each 2x2 block's rounded mean."""
    a = hr.astype(np.int64)
    s = a[0::2, 0::2] + a[0::2, 1::2] + a[1::2, 0::2] + a[1::2, 1::2]
    return ((s + 2) >> 2).astype(np.uint8)


def half_up(p: np.ndarray, s: int) -> np.ndarray:
    return p if s == 0 else (p + (1 << (s - 1))) >> s


def conv_int(x: np.ndarray, w: np.ndarray) -> np.ndarray:
    """Zero-padded convolution, stride 1: x [cin, h, w] int64, w [cout, cin, k, k] -> [cout, h, w]."""
    cout, cin, k, _ = w.shape
    p = k // 2
    _, h, ww = x.shape
    xp = np.zeros((cin, h + 2 * p, ww + 2 * p), np.int64)
    xp[:, p:p + h, p:p + ww] = x
    cols = np.empty((cin, k, k, h, ww), np.int64)
    for i in range(k):
        for j in range(k):
            cols[:, i, j] = xp[:, i:i + h, j:j + ww]
    return np.tensordot(w.astype(np.int64), cols, axes=([1, 2, 3], [0, 1, 2]))


def int_forward(prior: dict, lr: np.ndarray) -> np.ndarray:
    """`wai-upscale:sr2` (§3.2) of [h, w, 3] u8 under a parsed prior -> [2h, 2w, 3] u8."""
    x = lr.astype(np.int64).transpose(2, 0, 1) - 128
    n = len(prior["layers"])
    for i, l in enumerate(prior["layers"]):
        a = conv_int(x, l["w"]) + l["b"][:, None, None]
        v = half_up(a * l["m"][:, None, None], l["shift"])
        x = np.clip(v, 0, prior["act_hi"]) if i + 1 < n else v
    c, h, w = x.shape
    d = x.reshape(3, 2, 2, h, w).transpose(0, 3, 1, 4, 2).reshape(3, 2 * h, 2 * w)
    out = up2(lr).astype(np.int64).transpose(2, 0, 1) + d
    return np.clip(out, 0, 255).astype(np.uint8).transpose(1, 2, 0)


def write_wsr1(prior: dict) -> bytes:
    """The `WSR1` bytes of a prior (§3.1)."""
    b = bytearray(b"WSR1")
    b += struct.pack("<BBH", 1, len(prior["layers"]), prior["act_hi"])
    for l in prior["layers"]:
        cout, cin, k, _ = l["w"].shape
        b += struct.pack("<BBH", k, l["shift"], cout)
        b += l["w"].astype(np.int8).tobytes()
        b += l["b"].astype("<i4").tobytes()
        b += l["m"].astype("<u2").tobytes()
    return bytes(b)


def read_wsr1(data: bytes) -> dict:
    assert data[:4] == b"WSR1" and data[4] == 1
    n, act_hi = data[5], struct.unpack_from("<H", data, 6)[0]
    at, cin, layers = 8, 3, []
    for _ in range(n):
        k, shift, cout = struct.unpack_from("<BBH", data, at)
        at += 4
        nw = cout * cin * k * k
        w = np.frombuffer(data, np.int8, nw, at).astype(np.int64).reshape(cout, cin, k, k)
        at += nw
        b = np.frombuffer(data, "<i4", cout, at).astype(np.int64)
        at += 4 * cout
        m = np.frombuffer(data, "<u2", cout, at).astype(np.int64)
        at += 2 * cout
        layers.append({"w": w, "b": b, "m": m, "shift": shift})
        cin = cout
    assert at == len(data)
    return {"act_hi": act_hi, "layers": layers}


def psnr(a: np.ndarray, b: np.ndarray) -> float:
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    return float("inf") if mse == 0 else 10 * math.log10(255 * 255 / mse)


# ---------------------------------------------------------------------------------------
# The generated pictures.
# ---------------------------------------------------------------------------------------

SS = 4  # supersampling: every picture is rendered at 4x and box-filtered to its size


def render(seed: int, h: int, w: int, family: str) -> np.ndarray:
    """A picture of `family` from `seed`: `shapes` (filled polygons and ellipses over a
    gradient), `strokes` (rows of short strokes of several widths, like lettering),
    `texture` (gratings of several frequencies and orientations, masked into regions) or
    `mixed` (all three). Rendered at `SS`x and box-filtered, so edges are antialiased."""
    rng = np.random.default_rng(seed)
    H, W = h * SS, w * SS
    yy, xx = np.mgrid[0:H, 0:W].astype(np.float32)
    yy = (yy + 0.5) / SS
    xx = (xx + 0.5) / SS

    def colour():
        return rng.integers(0, 256, 3).astype(np.float32)

    th = rng.uniform(0, 2 * math.pi)
    t = (xx * math.cos(th) + yy * math.sin(th)) / max(h, w)
    t = (t - t.min()) / max(1e-6, float(t.max() - t.min()))
    c0, c1 = colour(), colour()
    img = c0[None, None, :] * (1 - t[..., None]) + c1[None, None, :] * t[..., None]

    def paint(mask, col):
        nonlocal img
        img = np.where(mask[..., None], col[None, None, :], img)

    def shapes(n):
        for _ in range(n):
            cx, cy = rng.uniform(0, w), rng.uniform(0, h)
            r = rng.uniform(0.04, 0.3) * min(h, w)
            if rng.random() < 0.5:
                a, b, phi = r, r * rng.uniform(0.3, 1.0), rng.uniform(0, math.pi)
                u = (xx - cx) * math.cos(phi) + (yy - cy) * math.sin(phi)
                v = -(xx - cx) * math.sin(phi) + (yy - cy) * math.cos(phi)
                paint((u / a) ** 2 + (v / b) ** 2 <= 1, colour())
            else:
                k = int(rng.integers(3, 7))
                ang = np.sort(rng.uniform(0, 2 * math.pi, k))
                px = cx + r * np.cos(ang)
                py = cy + r * np.sin(ang)
                inside = np.ones_like(xx, dtype=bool)
                for i in range(k):
                    x0, y0, x1, y1 = px[i], py[i], px[(i + 1) % k], py[(i + 1) % k]
                    inside &= (x1 - x0) * (yy - y0) - (y1 - y0) * (xx - x0) >= 0
                paint(inside, colour())

    def strokes(rows):
        ink = colour() * 0.3
        for r in range(rows):
            y = (r + 0.5) * h / rows + rng.uniform(-1, 1)
            x = rng.uniform(0, w * 0.1)
            width = rng.uniform(0.6, 2.5)
            size = rng.uniform(0.35, 0.8) * h / rows
            while x < w:
                for _ in range(int(rng.integers(1, 4))):
                    x0, y0 = x + rng.uniform(0, size), y + rng.uniform(-size / 2, size / 2)
                    x1, y1 = x + rng.uniform(0, size), y + rng.uniform(-size / 2, size / 2)
                    dx, dy = x1 - x0, y1 - y0
                    L2 = dx * dx + dy * dy + 1e-6
                    tt = np.clip(((xx - x0) * dx + (yy - y0) * dy) / L2, 0, 1)
                    d2 = (xx - x0 - tt * dx) ** 2 + (yy - y0 - tt * dy) ** 2
                    paint(d2 <= (width / 2) ** 2, ink)
                x += size * rng.uniform(1.1, 1.8)

    def texture(n):
        nonlocal img
        for _ in range(n):
            f = rng.uniform(0.03, 0.3)
            phi = rng.uniform(0, math.pi)
            g = 0.5 + 0.5 * np.sin(2 * math.pi * f * (xx * math.cos(phi) + yy * math.sin(phi)) + rng.uniform(0, 6.3))
            cx, cy, r = rng.uniform(0, w), rng.uniform(0, h), rng.uniform(0.15, 0.5) * min(h, w)
            m = ((xx - cx) ** 2 + (yy - cy) ** 2) <= r * r
            ca, cb = colour(), colour()
            tex = ca[None, None, :] * g[..., None] + cb[None, None, :] * (1 - g[..., None])
            img = np.where(m[..., None], tex, img)

    if family == "shapes":
        shapes(int(rng.integers(6, 14)))
    elif family == "strokes":
        shapes(int(rng.integers(1, 4)))
        strokes(int(rng.integers(3, 8)))
    elif family == "texture":
        texture(int(rng.integers(2, 5)))
        shapes(int(rng.integers(1, 4)))
    else:
        texture(int(rng.integers(1, 3)))
        shapes(int(rng.integers(3, 8)))
        strokes(int(rng.integers(2, 5)))
    out = img.reshape(h, SS, w, SS, 3).mean(axis=(1, 3))
    return np.clip(np.floor(out + 0.5), 0, 255).astype(np.uint8)


FAMILIES = ["shapes", "strokes", "texture", "mixed"]


def title_frames(seed: int, size: int, n: int, family: str) -> list[np.ndarray]:
    """A title: a scene twice the frame's size, panned under a window that moves a few
    pixels each frame."""
    canvas = render(seed, 2 * size, 2 * size, family)
    rng = np.random.default_rng(seed ^ 0x5EED)
    dy, dx = int(rng.integers(-9, 10)), int(rng.integers(-9, 10))
    y0, x0 = size // 2, size // 2
    frames = []
    for i in range(n):
        y = int(np.clip(y0 + i * dy, 0, size))
        x = int(np.clip(x0 + i * dx, 0, size))
        frames.append(np.ascontiguousarray(canvas[y:y + size, x:x + size]))
    return frames


# ---------------------------------------------------------------------------------------
# Training.
# ---------------------------------------------------------------------------------------


def train(args):
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    torch.manual_seed(args.seed)
    torch.set_num_threads(args.threads)
    out = Path(args.out)
    (out / "pictures").mkdir(parents=True, exist_ok=True)
    (out / "priors").mkdir(parents=True, exist_ok=True)
    manifest = {"pictures": [], "titles": {}}

    def keep(name, pic, group):
        (out / "pictures" / f"{name}.rgb").write_bytes(pic.tobytes())
        manifest["pictures"].append({"name": name, "group": group, "width": pic.shape[1], "height": pic.shape[0],
                                     "sha256": hashlib.sha256(pic.tobytes()).hexdigest()})

    train_set = [render(args.seed * 1000 + i, args.train_size, args.train_size, FAMILIES[i % 4]) for i in range(args.train_n)]
    held = []
    for i in range(args.held_n):
        p = render(args.seed * 1000 + 500 + i, args.size, args.size, FAMILIES[i % 4])
        held.append((f"held-{FAMILIES[i % 4]}-{i}", p))
        keep(held[-1][0], p, "held")
    titles = {}
    for i, fam in enumerate(FAMILIES):
        name = f"title-{fam}"
        titles[name] = title_frames(args.seed * 1000 + 900 + i, args.size, args.title_frames, fam)
    if args.source:
        from PIL import Image
        titles["title-source"] = [np.asarray(Image.open(args.source).convert("RGB"), np.uint8)]
    for name, frames in titles.items():
        manifest["titles"][name] = []
        for j, f in enumerate(frames):
            keep(f"{name}-{j}", f, name)
            manifest["titles"][name].append(f"{name}-{j}")
    (out / "pictures" / "manifest.json").write_text(json.dumps(manifest, indent=1) + "\n")

    act_c = float(args.act_ceiling)
    s_h = 1.0 / (1 << args.act_frac)

    class SR(nn.Module):
        def __init__(self):
            super().__init__()
            chans = [3] + [args.width] * (args.layers - 1) + [12]
            self.convs = nn.ModuleList([nn.Conv2d(chans[i], chans[i + 1], 3, padding=1) for i in range(args.layers)])
            nn.init.zeros_(self.convs[-1].weight)
            nn.init.zeros_(self.convs[-1].bias)
            self.quant = False

        def fq_w(self, w):
            if not self.quant:
                return w
            s = w.detach().abs().amax(dim=(1, 2, 3), keepdim=True).clamp_min(1e-8) / 127
            return w + (torch.round(w / s).clamp(-127, 127) * s - w).detach()

        def forward(self, lr, base):
            x = lr - 128
            for i, c in enumerate(self.convs):
                x = F.conv2d(x, self.fq_w(c.weight), c.bias, padding=1)
                if i + 1 < len(self.convs):
                    x = x.clamp(0, act_c)
                    if self.quant:
                        x = x + (torch.round(x / s_h) * s_h - x).detach()
            return base + F.pixel_shuffle(x, 2)

    def batch_from(pics, rng, n, patch):
        lrs, bases, hrs = [], [], []
        for _ in range(n):
            p = pics[int(rng.integers(len(pics)))]
            y = int(rng.integers(0, (p.shape[0] - patch) // 2 + 1)) * 2
            x = int(rng.integers(0, (p.shape[1] - patch) // 2 + 1)) * 2
            hr = p[y:y + patch, x:x + patch]
            if rng.random() < 0.5:
                hr = hr[:, ::-1]
            if rng.random() < 0.5:
                hr = hr[::-1]
            if rng.random() < 0.5:
                hr = hr.transpose(1, 0, 2)
            hr = np.ascontiguousarray(hr)
            lr = box_down(hr)
            lrs.append(lr)
            bases.append(up2(lr))
            hrs.append(hr)
        t = lambda a: torch.from_numpy(np.stack(a).astype(np.float32)).permute(0, 3, 1, 2)
        return t(lrs), t(bases), t(hrs)

    def fit(model, pics, steps, lr0, seed, quant_from):
        rng = np.random.default_rng(seed)
        opt = torch.optim.Adam(model.parameters(), lr=lr0)
        losses = []
        for step in range(steps):
            model.quant = step >= quant_from
            for g in opt.param_groups:
                g["lr"] = lr0 * (0.5 * (1 + math.cos(math.pi * step / steps)))
            lr_t, base_t, hr_t = batch_from(pics, rng, args.batch, args.patch)
            loss = torch.sqrt((model(lr_t, base_t) - hr_t) ** 2 + 1e-3).mean()  # Charbonnier: residual bytes follow |error|
            opt.zero_grad()
            loss.backward()
            opt.step()
            if step % 250 == 0 or step + 1 == steps:
                losses.append([step, round(float(loss), 5)])
        return losses

    def quantise(model) -> dict:
        layers = []
        s_in = 1.0
        n = len(model.convs)
        for i, c in enumerate(model.convs):
            w = c.weight.detach().double().numpy()
            b = c.bias.detach().double().numpy()
            sw = np.maximum(np.abs(w).max(axis=(1, 2, 3)), 1e-8) / 127
            wq = np.clip(np.round(w / sw[:, None, None, None]), -127, 127).astype(np.int64)
            s_out = s_h if i + 1 < n else 1.0
            ratio = sw * s_in / s_out
            shift = 31
            while shift > 0 and np.round(ratio.max() * (1 << shift)) > 65535:
                shift -= 1
            m = np.clip(np.round(ratio * (1 << shift)), 1, 65535).astype(np.int64)
            bq = np.round(b / (sw * s_in)).astype(np.int64)
            layers.append({"w": wq, "b": bq, "m": m, "shift": shift})
            s_in = s_out
        return {"act_hi": int(round(act_c / s_h)), "layers": layers}

    def evaluate(prior, pairs):
        rows = []
        for name, hr in pairs:
            lr = box_down(hr)
            rows.append({"picture": name, "up2_psnr": round(psnr(up2(lr), hr), 3), "sr_psnr": round(psnr(int_forward(prior, lr), hr), 3)})
        return rows

    report = {"seed": args.seed, "width": args.width, "layers": args.layers, "act_ceiling": args.act_ceiling, "act_frac": args.act_frac,
              "train": {"pictures": args.train_n, "size": args.train_size, "steps": args.steps, "quant_from": args.steps * 3 // 4,
                        "patch": args.patch, "batch": args.batch, "lr": args.lr}}
    generic = SR()
    report["train"]["losses"] = fit(generic, train_set, args.steps, args.lr, args.seed, args.steps * 3 // 4)
    gq = quantise(generic)
    gbytes = write_wsr1(gq)
    assert write_wsr1(read_wsr1(gbytes)) == gbytes
    (out / "priors" / "generic.wsr").write_bytes(gbytes)
    report["generic"] = {"bytes": len(gbytes), "sha256": hashlib.sha256(gbytes).hexdigest(), "held": evaluate(gq, held)}
    report["titles"] = {}
    for k, (name, frames) in enumerate(titles.items()):
        model = SR()
        model.load_state_dict(generic.state_dict())
        losses = fit(model, frames, args.title_steps, args.lr * 0.3, args.seed * 7919 + k, 0)
        tq = quantise(model)
        tb = write_wsr1(tq)
        (out / "priors" / f"{name}.wsr").write_bytes(tb)
        pairs = [(f"{name}-{j}", f) for j, f in enumerate(frames)]
        report["titles"][name] = {"bytes": len(tb), "sha256": hashlib.sha256(tb).hexdigest(), "losses": losses,
                                  "generic": evaluate(gq, pairs), "title": evaluate(tq, pairs)}
    (out / "report.json").write_text(json.dumps(report, indent=1) + "\n")
    print(json.dumps({"generic": report["generic"]["bytes"], "held": report["generic"]["held"]}, indent=1))
    for name, t in report["titles"].items():
        g = np.mean([r["sr_psnr"] for r in t["generic"]])
        tt = np.mean([r["sr_psnr"] for r in t["title"]])
        u = np.mean([r["up2_psnr"] for r in t["generic"]])
        print(f"{name}: up2 {u:.2f} dB, generic {g:.2f} dB, per-title {tt:.2f} dB, prior {t['bytes']} B")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--out", required=True)
    ap.add_argument("--seed", type=int, default=20261021)
    ap.add_argument("--width", type=int, default=16)
    ap.add_argument("--layers", type=int, default=3)
    ap.add_argument("--steps", type=int, default=4000)
    ap.add_argument("--title-steps", type=int, default=1500)
    ap.add_argument("--lr", type=float, default=2e-3)
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--patch", type=int, default=48)
    ap.add_argument("--train-n", type=int, default=240)
    ap.add_argument("--train-size", type=int, default=96)
    ap.add_argument("--held-n", type=int, default=8)
    ap.add_argument("--size", type=int, default=128)
    ap.add_argument("--title-frames", type=int, default=8)
    ap.add_argument("--act-ceiling", type=float, default=255.0)
    ap.add_argument("--act-frac", type=int, default=3)
    ap.add_argument("--threads", type=int, default=8)
    ap.add_argument("--source", help="a PNG to add as one more title (the measurement passes the pinned source crop)")
    train(ap.parse_args())


if __name__ == "__main__":
    main()