#!/usr/bin/env python3
"""Build every compressed variant BEFORE any model run.
  --probe            : bpp of each codec over its parameter grid (used only to pick rate targets; reads no model output)
  --generate T1,T2   : per image, per codec, pick the parameter whose bpp is closest to each target (log-space), decode to PNG.
bpp = bitstream bits / (original width x height)  -- always at ORIGINAL resolution, never at the 336px model input.
Decoded PNG (lossless) is what the model sees, for every arm including the original."""
import argparse, io, json, math, os, subprocess, sys, tempfile, time, hashlib
from pathlib import Path
import numpy as np, torch
from PIL import Image
PKG = Path(__file__).resolve().parents[2]; D = PKG / "work/data"
ap = argparse.ArgumentParser(); ap.add_argument("--probe", action="store_true"); ap.add_argument("--generate")
ap.add_argument("--control", default="jpeg_q2"); ap.add_argument("--bpgenc", required=True); ap.add_argument("--bpgdec", required=True)
ap.add_argument("--png_out", required=True, help="dir for decoded PNGs (large; kept outside the package)")
a = ap.parse_args()
torch.set_num_threads(2)
sel = json.load(open(D / "selection.json")); names = [n for s in sel["subsets"].values() for n in s]
def load(n): return Image.open(D / "img_orig" / n).convert("RGB")
def bpp(nbytes, im): return nbytes * 8 / (im.width * im.height)
def enc_pil(im, fmt, q):
    b = io.BytesIO(); im.save(b, fmt, quality=q, **({"method": 6} if fmt == "WEBP" else {})); data = b.getvalue()
    return data, Image.open(io.BytesIO(data)).convert("RGB")
TMP = Path(tempfile.mkdtemp(prefix="bpg_", dir=a.png_out))
def enc_bpg(im, n, q):
    src = TMP / "src.png"; im.save(src); out = TMP / "x.bpg"; dec = TMP / "x.png"
    subprocess.run([a.bpgenc, "-q", str(q), "-o", str(out), str(src)], check=True, capture_output=True)
    subprocess.run([a.bpgdec, "-o", str(dec), str(out)], check=True, capture_output=True)
    return out.read_bytes(), Image.open(dec).convert("RGB")
from compressai.zoo import mbt2018_mean
from torchvision.transforms.functional import to_tensor, to_pil_image
NETS = {}
def enc_mbt(im, q):
    if q not in NETS:
        NETS[q] = mbt2018_mean(quality=q, pretrained=True).eval(); NETS[q].update()
    x = to_tensor(im).unsqueeze(0); h, w = x.shape[-2:]; H = (h + 63) // 64 * 64; W = (w + 63) // 64 * 64
    xp = torch.nn.functional.pad(x, (0, W - w, 0, H - h), mode="replicate")
    with torch.no_grad():
        o = NETS[q].compress(xp); r = NETS[q].decompress(o["strings"], o["shape"])["x_hat"].clamp(0, 1)[..., :h, :w]
    nbytes = sum(len(s[0]) for s in o["strings"]) + 4  # +4 bytes to store the latent shape header
    return nbytes, to_pil_image(r[0])
GRIDS = {"jpeg": list(range(1, 96)), "webp": list(range(0, 101)), "bpg": list(range(51, -1, -1)), "mbt2018mean": list(range(1, 9))}
def encode(codec, im, n, p):
    if codec == "jpeg": d, r = enc_pil(im, "JPEG", p); return len(d), r
    if codec == "webp": d, r = enc_pil(im, "WEBP", p); return len(d), r
    if codec == "bpg": d, r = enc_bpg(im, n, p); return len(d), r
    return enc_mbt(im, p)
if a.probe:
    probe = {}
    for codec, grid in GRIDS.items():
        sub = grid if codec == "mbt2018mean" else grid[::max(1, len(grid) // 12)]
        probe[codec] = {str(p): [round(bpp(encode(codec, load(n), n, p)[0], load(n)), 4) for n in names[::6]] for p in sub}
        print(codec, {p: round(float(np.median(v)), 3) for p, v in probe[codec].items()}, flush=True)
    json.dump(probe, open(D / "codec_probe.json", "w"), indent=1); sys.exit()
targets = [float(t) for t in a.generate.split(",")]
out = Path(a.png_out); rows = []
def save(arm, n, rec, nbytes, param, im):
    f = out / arm / (n[:-4] + ".png"); f.parent.mkdir(parents=True, exist_ok=True); rec.save(f)
    arr0 = np.asarray(im, dtype=np.float64); arr1 = np.asarray(rec, dtype=np.float64); mse = ((arr0 - arr1) ** 2).mean()
    rows.append(dict(arm=arm, image=n, param=param, bytes=nbytes, bpp=round(bpp(nbytes, im), 5), width=im.width, height=im.height,
                     psnr_vs_coco_jpeg=round(10 * math.log10(255 ** 2 / mse), 3) if mse > 0 else None, png_sha256=hashlib.sha256(f.read_bytes()).hexdigest()))
t0 = time.time()
for i, n in enumerate(names):
    im = load(n)
    save("orig", n, im, (D / "img_orig" / n).stat().st_size, "coco_as_distributed", im)
    cq = int(a.control.split("_q")[1]); nb, r = encode(a.control.split("_q")[0], im, n, cq); save(a.control, n, r, nb, cq, im)
    for codec, grid in GRIDS.items():
        for ti, T in enumerate(targets):
            # monotone parameter grid -> binary search on log|bpp - T|, then check neighbours
            lo, hi, cache = 0, len(grid) - 1, {}
            def f(k):
                if k not in cache: cache[k] = encode(codec, im, n, grid[k])
                return cache[k]
            while hi - lo > 1:
                mid = (lo + hi) // 2
                if bpp(f(mid)[0], im) < T: lo = mid
                else: hi = mid
            best = min((lo, hi), key=lambda k: abs(math.log(bpp(f(k)[0], im) / T)))
            nb, r = f(best); save(f"{codec}_t{ti+1}", n, r, nb, grid[best], im)
    print(f"[{i+1}/{len(names)}] {n} {time.time()-t0:.0f}s", flush=True)
json.dump(dict(targets_bpp=targets, control=a.control, rows=rows), open(D / "variants.json", "w"), indent=0)
print("done", len(rows), "variants")