#!/usr/bin/env python3
"""Run a trained NadirNet over held-out pages and write render jobs (Python port of nadir-infer.js decode)."""
import argparse, gzip, json, math, pickle
from pathlib import Path
import numpy as np
import torch

from model import NadirNet
from train import collate, prep
from mkjobs import job_from

TREAT = ["keep", "invert", "adapt", "remove"]
RENDER_T = {"keep": "keep", "invert": "invert", "adapt": "invert", "remove": "clear"}  # render.py has no adapt path


from safety import oklab, lab_to_rgb, contrast, safe_treatment, enforce_contrast, is_media_keep, hue_guard, keep_invisible, rescue_surfaces, glass_panels, opaque_surfaces, dark_washes


def decode(L, lab, alpha_logit, img_logits, rels):
    toks = L["tokens"][:256]
    out, keep = [], set()
    for i, t in enumerate(toks):
        ain = min(1 - 1e-3, max(1e-3, t["rgba"][3]))
        alpha = 1 / (1 + math.exp(-(alpha_logit[i] + math.log(ain / (1 - ain)))))
        a = 1.0 if (t["rgba"][3] >= 0.999 and alpha > 0.97) else alpha
        out.append({"rgb": lab_to_rgb(*lab[i]), "a": a, "L": float(lab[i][0])})
        if is_media_keep(t):  # shown as-is: later safety passes see its real colour
            keep.add(i)
            rgb = [int(x) for x in t["rgba"][:3]]
            out[i] = {"rgb": rgb, "a": float(t["rgba"][3]), "L": oklab(*rgb)[0]}
        else:
            hue_guard(t, out[i])
    glass_panels(toks, out, keep)
    for i, t in enumerate(toks):
        if t["f"][24] > 0.5 and t["role"] == 0 and out[i]["L"] > 0.35:
            o = oklab(*out[i]["rgb"]); out[i]["rgb"] = lab_to_rgb(0.2, o[1] * 0.3, o[2] * 0.3); out[i]["L"] = 0.2
    opaque_surfaces(toks, out, keep)
    dark_washes(toks, out, keep)
    keep_invisible(toks, out, rels, keep)
    enforce_contrast(toks, out, rels, keep)
    rescue_surfaces(toks, out, rels, keep)
    def hx(o):
        h = "#%02x%02x%02x" % tuple(o["rgb"])
        return h + ("%02x" % round(o["a"] * 255) if o["a"] < 0.999 else "")
    colors = [None if i in keep else hx(o) for i, o in enumerate(out)]
    treat = {}
    for j, im in enumerate(L["images"][:32]):
        t = TREAT[int(np.argmax(img_logits[j]))] if j < len(img_logits) else "keep"
        treat[str(j)] = RENDER_T[t]
    return colors, sorted(keep), treat


def element_treatments(model, s, lab, ctx, dev):
    """Element head over a page's first 32 images (render jobs only know those)."""
    m = min(32, len(s["IF"]))
    if not m:
        return []
    f = torch.from_numpy(s["IF"][:m])[None].to(dev)
    th = torch.from_numpy(s["thumbs"][:m].transpose(0, 3, 1, 2) / 255.0).float()[None].to(dev)
    ul = torch.from_numpy(s["IF"][:m, 23:26] * np.array([1, 0.25, 0.25], np.float32))[None].to(dev)
    u = np.where((s["img_under"][:m] >= 0) & (s["img_under"][:m] < lab.shape[1]), s["img_under"][:m], s["root"])
    ud = lab[:, torch.from_numpy(u).to(dev)]
    t, _ = model.elements(f, th, ul, ud, ctx)
    return t[0].float().cpu().numpy().tolist()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--crawl", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--n", type=int, default=150)
    ap.add_argument("--src", type=int, default=0)
    ap.add_argument("--tag", default="nadir")
    ap.add_argument("--also", default="haiku,dr", help="extra render tags to emit for the same pages")
    ap.add_argument("--haiku", default="")
    args = ap.parse_args()
    ck = torch.load(args.ckpt, map_location="cpu")
    model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(dev)
    data = pickle.load(open(args.data, "rb"))
    val = [s for s in data["val"] if s["src"] == args.src][: args.n]
    with open(args.out, "w") as f:
        for s in val:
            s = prep(s)
            L = json.load(gzip.open(Path(args.crawl) / "pages" / s["id"] / "light.json.gz"))
            b = collate([s], dev)
            with torch.no_grad():
                lab, al, ctx = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], torch.zeros_like(b["src"]), b["rel"], b["mask"])
                il = element_treatments(model, s, lab.float(), ctx, dev)
            n = len(s["role"])
            colors, keep, treat = decode(L, lab[0, :n].float().cpu().numpy().tolist(), al[0, :n].float().cpu().numpy().tolist(),
                                         il, s["rel"])
            colors = colors + [None] * (len(L["tokens"]) - len(colors))
            j = job_from(L, s["id"], args.tag, colors, treat)
            j["keep"] = sorted(set(j["keep"]) | set(keep))
            f.write(json.dumps(j) + "\n")
            if "dr" in args.also:
                f.write(json.dumps(job_from(L, s["id"], "dr", [None] * len(L["tokens"]), {}, baseline=True)) + "\n")
            if "haiku" in args.also and args.haiku:
                hp = Path(args.haiku) / f"{s['id']}.json"
                if hp.exists():
                    h = json.loads(hp.read_text())
                    outs = [h["colors"].get(str(i)) for i in range(len(L["tokens"]))]
                    f.write(json.dumps(job_from(L, s["id"], "haiku", outs, h["images"])) + "\n")
    print("wrote", len(val), "pages")


if __name__ == "__main__":
    main()