"""Turn ink_9um outputs into the page's ZeroGPU layer.

Input: gpu/<winding>/receipt.json plus the forward TIFF it names (uint8, one per winding, the 9.362 um
surface-volume grid). Those TIFFs (about 25 MB each) are not shipped with the Space; anyone can regenerate
them with the model card's command on the public surface volumes listed in data/gpu/receipts/.
Writes data/gpu/<w>.jpg (3D texture, display-filtered), data/hi/<w>_gpu.jpg (close-up view, the model card's
display rescale only), data/gpu/receipts/<w>.json (the run records) and data/gpu.json (the page's summary).
UNROLL_GPU and UNROLL_OUT override the input and output folders.
"""
import glob
import json
import os
import sys

import numpy as np
import tifffile
from PIL import Image

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, f"{ROOT}/tools")
from build import CELL_MM, HI_MAX, PPMM, local_contrast  # noqa: E402

OUT = os.environ.get("UNROLL_OUT", f"{ROOT}/data")
GPU_IN = os.environ.get("UNROLL_GPU", f"{ROOT}/gpu")
# run-record fields that are local plumbing, not provenance
DROP = {"space", "status"}


def say(items):
    """'a', 'a and b', 'a, b and c'."""
    return items[0] if len(items) == 1 else f"{', '.join(items[:-1])} and {items[-1]}"


def main():
    os.makedirs(f"{OUT}/gpu/receipts", exist_ok=True)
    manifest = json.load(open(f"{OUT}/scroll.json"))
    by_id = {w["id"]: w for w in manifest["windings"]}
    wins, tot_patches, tot_measured, tot_calls = {}, 0, 0.0, 0
    device = None
    dirs = sorted(glob.glob(f"{GPU_IN}/w*/"))
    if not dirs:
        sys.exit(f"no ZeroGPU outputs under {GPU_IN}/w*/ (see the docstring)")
    for d in dirs:
        w = os.path.basename(d.rstrip("/"))
        receipt = json.load(open(f"{d}receipt.json"))
        record = {k: v for k, v in receipt.items() if k not in DROP}
        record["gpu_seconds"] = {k: v for k, v in receipt["gpu_seconds"].items() if k != "note"}
        record["inputs"] = {"direction": receipt["inputs"]["direction"], "geometry": receipt["inputs"]["geometry"],
                            "source": receipt["inputs"]["source"]}
        json.dump(record, open(f"{OUT}/gpu/receipts/{w}.json", "w"), indent=1, sort_keys=True)
        tif = f"{d}{receipt['outputs']['forward']['file']}"
        a = tifffile.imread(tif).astype(np.float32) / 255
        win = by_id[w]
        assert a.shape == (win["gridH"] * 20, win["gridW"] * 20), (w, a.shape)
        disp = np.clip((a - 0.25) / 0.5, 0, 1)          # model card: display rescale (p - 0.25) / 0.5
        disp[a == 0] = 0
        im = Image.fromarray((disp * 255).round().astype(np.uint8))
        hi = im.copy()
        hi.thumbnail((HI_MAX, HI_MAX), Image.LANCZOS)
        hi.save(f"{OUT}/hi/{w}_gpu.jpg", quality=84, optimize=True, progressive=True)
        tw = int(round(win["gridW"] * CELL_MM * PPMM * 1.5))
        th = int(round(win["gridH"] * CELL_MM * PPMM * 1.5))
        tex = local_contrast(im.resize((tw, th), Image.LANCZOS), 2.2 * PPMM * 1.5)
        tex.save(f"{OUT}/gpu/{w}.jpg", quality=86, optimize=True)
        g = receipt["gpu_seconds"]
        p = receipt["patches"]["forward"]["nonempty"]
        device = (receipt.get("device") or {}).get("device_name")
        tot_patches += p
        tot_measured += g["measured"]
        tot_calls += g["calls"]
        args = receipt["args"]
        wins[w] = {
            "tex": f"data/gpu/{w}.jpg",
            "receipt": f"data/gpu/receipts/{w}.json",
            "hi": f"hi/{w}_gpu.jpg",
            "patches": p,
            "gpuSeconds": round(g["measured"], 1),
            "calls": g["calls"],
            "overlap": args["overlap"],
            "started": receipt["calls"][0]["started"],
            "meta": (f"ink_9um (9 µm): scrollprize/ink_9um (seed 42, step 75000) on the 9.4 µm scan (9.362 µm voxels) · "
                     f"{p:,} patches of 128×128 px, overlap {args['overlap']}"
                     f"{' (patch seams visible)' if args['overlap'] < 0.5 else ''} · {g['measured']:.1f} s of GPU in "
                     f"{g['calls']} call{'s' if g['calls'] > 1 else ''} · {receipt['calls'][0]['started'][:10]} · "
                     f"drag to pan, scroll or pinch to zoom, double-click to fit"),
        }
        print(w, p, g["measured"], tex.size)
    ids = sorted(wins)
    full = [w for w in ids if wins[w]["overlap"] >= 0.5]
    low = [w for w in ids if wins[w]["overlap"] < 0.5]
    out = {
        "windings": wins,
        "totals": {"windings": len(wins), "patches": tot_patches, "gpuSeconds": round(tot_measured, 1),
                   "calls": tot_calls, "device": device},
        "about": (f"<b>ink_9um (9 µm)</b> is the cyan band. Chase Hendrick ran the public "
                  f"<a href=\"https://huggingface.co/scrollprize/ink_9um\" target=\"_blank\" rel=\"noopener\">"
                  f"scrollprize/ink_9um</a> model, which the Vesuvius Challenge trained, over windings {ids[0]} to {ids[-1]} "
                  f"and processed the output for this viewer: {tot_patches:,} patches in "
                  f"{tot_measured:.0f} s of measured GPU time on half an NVIDIA RTX PRO 6000 Blackwell. It reads the "
                  f"9.4 µm scan (9.362 µm voxels), about four times coarser than the 2.4 µm scan behind the amber layer, so its ink map "
                  f"is noisier. {say(full)} ran with patch overlap 0.5; {say(low)} ran with less overlap, so "
                  f"their 128 px patch grid shows as faint square tiles. That grid is a stitching artifact, not ink. "
                  f"Click a cyan winding to compare the two maps."),
    }
    json.dump(out, open(f"{OUT}/gpu.json", "w"), indent=1)
    print(json.dumps(out["totals"]))


if __name__ == "__main__":
    main()