"""Build the Unroll page's data from public Vesuvius Challenge files (run tools/fetch.py first).

  python tools/fetch.py && python tools/build.py     (from the Space's root; needs numpy, scipy, tifffile, pillow)

Inputs come from raw/ (written by fetch.py); outputs go to data/. UNROLL_RAW and UNROLL_OUT override both.
Outputs:
  scroll.bin   int16 vertex positions (scan voxels minus the axis center) for every winding's decimated grid
  scroll.json  manifest: per-winding grid sizes, byte offsets, physical sizes, atlas rects, sources
  atlas_ink.jpg, atlas_ct.jpg   every winding packed at one physical scale (grayscale)
  hi/wNNN.jpg  per-winding Vesuvius Challenge ink preview for the close-up view

Coarseness: UNROLL_STEP sets the grid step of a full build (default 3). A lighter copy of a built mesh, for phones
and low-power devices, needs no raw downloads:
  python tools/build.py --coarsen 2 data/PHerc0172     writes scroll_lite.bin, scroll_lite.json, scroll_lite.i16p.gz
It keeps every 2nd sampled row and column (plus the last ones), so the UV mapping (rowIdx, colIdx over the same
gridH, gridW), the atlases, the ink layers and the flat layout are unchanged.
"""
import glob
import http.client
import subprocess
import time
import urllib.error
import json
import os
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor

import numpy as np
from scipy.ndimage import gaussian_filter
import tifffile
from PIL import Image

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAW = os.environ.get("UNROLL_RAW", f"{ROOT}/raw")
OUT = os.environ.get("UNROLL_OUT", f"{ROOT}/data")
B = "https://vesuvius-challenge-open-data.s3.amazonaws.com"
SCROLL = os.environ.get("UNROLL_SCROLL", "PHerc0139")
VOX_MM = float(os.environ.get("UNROLL_VOX_MM", "0.009362"))   # scan voxel of the meshes (9.362 um for PHerc0139)
CELL_MM = 20 * VOX_MM      # tifxyz scale 0.05: one grid cell is 20 voxels
CT_RES = os.environ.get("UNROLL_CT_RES", "9.362um")          # surface volume for the X-ray layer (any scan: it is a UV texture)
NO_INK = os.environ.get("UNROLL_NO_INK") == "1"
RENUMBER = os.environ.get("UNROLL_RENUMBER") == "1"          # name the windings w001.. by radius, innermost first
STEP = int(os.environ.get("UNROLL_STEP", "3"))   # grid decimation for the browser mesh (larger is coarser)
PPMM = float(os.environ.get("UNROLL_PPMM", "7.5"))   # atlas pixels per mm (lowered for long scrolls)
ATLAS = 4096
PAD = 6
HI_MAX = 2400
CT_LAYERS = (11, 18)       # surface-volume layers averaged for the X-ray texture

Image.MAX_IMAGE_PIXELS = None


def fetch_ct(w, src, shape_hw):
    """Level-3 X-ray texture (mean of the middle layers) of a winding's 9.362 um surface volume."""
    cache = f"{RAW}/{src.get('w', w)}/ct_l3.npy"
    if os.path.exists(cache):
        return np.load(cache)
    seg = src["seg"]
    sv = [s for s in src["sv"] if s.startswith(CT_RES) and "-L1" not in s and s.endswith(".zarr")][0]
    root = f"{B}/{SCROLL}/segments/{seg}/surface-volumes/{sv}"
    # the level nearest the grid's size (level 3 for 9 um scans; finer scans need a coarser level)
    lvl = 3
    for lv in range(3, 7):
        try:
            mt = json.loads(urllib.request.urlopen(f"{root}/{lv}/.zarray").read())
        except urllib.error.HTTPError:
            break
        lvl = lv
        if mt["shape"][1] <= shape_hw[0] * 8:
            break
    base = f"{root}/{lvl}"
    meta = json.loads(urllib.request.urlopen(f"{base}/.zarray").read())
    comp = meta.get("compressor")
    assert comp is None or comp.get("id") == "blosc", f"compressed zarr not handled: {comp}"
    D, H, W = meta["shape"]
    cz, ch, cw = meta["chunks"]
    sep = meta.get("dimension_separator") or "."
    out = np.zeros((H, W), np.float32)
    jobs = [(i, j) for i in range(-(-H // ch)) for j in range(-(-W // cw))]

    def one(ij):
        i, j = ij
        url = f"{base}/0{sep}{i}{sep}{j}"
        for attempt in range(6):
            r = subprocess.run(["curl", "-s", "-w", "%{http_code}", "-o", "-", url], capture_output=True)
            code = r.stdout[-3:]
            body = r.stdout[:-3]
            if code in (b"403", b"404"):
                return ij, None
            if code == b"200" and comp:
                import imagecodecs
                body = imagecodecs.blosc_decode(body)
            if code == b"200" and len(body) == cz * ch * cw:
                raw = body
                break
            time.sleep(1 + attempt)
        else:
            raise OSError(f"could not fetch {url}")
        a = np.frombuffer(raw, np.uint8).reshape(cz, ch, cw)
        lo = CT_LAYERS[0] if cz >= 28 else max(0, cz // 2 - 3)
        return ij, a[lo:lo + (CT_LAYERS[1] - CT_LAYERS[0])].astype(np.float32).mean(0)

    with ThreadPoolExecutor(4) as ex:
        for (i, j), tile in ex.map(one, jobs):
            if tile is None:
                continue
            h = min(ch, H - i * ch)
            ww = min(cw, W - j * cw)
            out[i * ch:i * ch + h, j * cw:j * cw + ww] = tile[:h, :ww]
    np.save(cache, out)
    return out


def axis_fit(grids):
    """Scroll axis as x(z), y(z) lines through the innermost winding's ring centroids."""
    x, y, z, v = grids[0]
    zs, cx, cy = [], [], []
    for z0 in np.arange(np.percentile(z[v], 5), np.percentile(z[v], 95), 150):
        sel = v & (np.abs(z - z0) < 40)
        if sel.sum() > 50:
            zs.append(z0)
            cx.append(x[sel].mean())
            cy.append(y[sel].mean())
    zs, cx, cy = map(np.array, (zs, cx, cy))
    ax = np.polyfit(zs, cx, 1)
    ay = np.polyfit(zs, cy, 1)
    return ax, ay


def local_contrast(img, sigma_px, base=0.16, gain=1.7):
    """Display filter for the 3D glow: subtract the local mean so letter-scale strokes stand out and
    broad bright patches (common in the inner windings' predictions) stop flooding the view."""
    a = np.asarray(img, np.float32) / 255
    m = (a > 0.03).astype(np.float32)
    mean = gaussian_filter(a * m, sigma_px) / np.maximum(gaussian_filter(m, sigma_px), 1e-3)
    out = np.clip(base + (a - mean) * gain, 0, 1) * m
    return Image.fromarray((out * 255).round().astype(np.uint8))


def shelf_pack(sizes, width):
    """Shelf packing, tallest first. sizes: {key: (w, h)} -> {key: (x, y)}, total height."""
    order = sorted(sizes, key=lambda k: -sizes[k][1])
    pos, x, y, shelf = {}, PAD, PAD, 0
    for k in order:
        w, h = sizes[k]
        if x + w + PAD > width:
            x, y, shelf = PAD, y + shelf + PAD, 0
        pos[k] = (x, y)
        x += w + PAD
        shelf = max(shelf, h)
    return pos, y + shelf + PAD


def main():
    os.makedirs(f"{OUT}/hi", exist_ok=True)
    dirs = sorted(glob.glob(f"{RAW}/[ws]*[0-9]"))
    names = [os.path.basename(d) for d in dirs]
    grids, srcs, metas = [], [], []
    for d in dirs:
        x, y, z = (tifffile.imread(f"{d}/{c}.tif").astype(np.float32) for c in "xyz")
        v = (x > 0) & (y > 0) & (z > 0)
        grids.append((x, y, z, v))
        srcs.append(json.load(open(f"{d}/source.json")))
        metas.append(json.load(open(f"{d}/meta.json")))
    if RENUMBER:
        # order by median distance from a first axis fit through all windings: the innermost becomes w001
        allg = [np.concatenate([g[i][g[3]] for g in grids]) for i in range(3)]
        ax0, ay0 = np.polyfit(allg[2], allg[0], 1), np.polyfit(allg[2], allg[1], 1)
        rad = [float(np.median(np.hypot(g[0][g[3]] - np.polyval(ax0, g[2][g[3]]), g[1][g[3]] - np.polyval(ay0, g[2][g[3]]))))
               for g in grids]
        order = np.argsort(rad)
        grids, srcs, metas, dirs = ([L[i] for i in order] for L in (grids, srcs, metas, dirs))
        names = [f"w{k + 1:03d}" for k in range(len(dirs))]
        print("renumbered by radius:", [(n, srcs[k]["seg"], round(rad[order[k]])) for k, n in enumerate(names)])
    ax, ay = axis_fit(grids)
    rad2 = [float(np.median(np.hypot(g[0][g[3]] - np.polyval(ax, g[2][g[3]]), g[1][g[3]] - np.polyval(ay, g[2][g[3]])))) * VOX_MM
            for g in grids]
    if RENUMBER:
        # second pass on the refined axis, so the ids follow the radius exactly
        order = np.argsort(rad2)
        grids, srcs, metas, dirs = ([L[i] for i in order] for L in (grids, srcs, metas, dirs))
        rad2 = [rad2[i] for i in order]
        ax, ay = axis_fit(grids)
    print("winding order check (id, median radius mm):", [(n, round(r, 1)) for n, r in zip(names, rad2)])
    bad = sum(1 for a, b in zip(rad2, rad2[1:]) if b < a)
    print("radius decreases between consecutive ids:", bad)
    global PPMM
    if "UNROLL_PPMM" not in os.environ:
        cells = sum(g[3].shape[0] * g[3].shape[1] for g in grids) * CELL_MM * CELL_MM
        PPMM = min(PPMM, float(np.sqrt(0.62 * ATLAS * ATLAS / cells)))
        print("PPMM", round(PPMM, 3))
    allz = np.concatenate([g[2][g[3]] for g in grids])
    zc = float((allz.min() + allz.max()) / 2)
    cx0, cy0 = float(np.polyval(ax, zc)), float(np.polyval(ay, zc))
    print("axis at z", round(zc), "center", round(cx0), round(cy0), "tilt", ax[0], ay[0])

    # textures at one physical scale
    sizes, ink_tiles, ct_tiles = {}, {}, {}
    for w, d, (x, y, z, v), src in zip(names, dirs, grids, srcs):
        H, W = v.shape
        tw, th = int(round(W * CELL_MM * PPMM)), int(round(H * CELL_MM * PPMM))
        sizes[w] = (tw, th)
        if not NO_INK:
            ink = Image.open(f"{d}/ink2um_ds8.jpg").convert("L")
            ink_tiles[w] = local_contrast(ink.resize((tw, th), Image.LANCZOS), 2.2 * PPMM)
            hi = ink.copy()
            hi.thumbnail((HI_MAX, HI_MAX), Image.LANCZOS)
            hi.save(f"{OUT}/hi/{w}.jpg", quality=84, optimize=True, progressive=True)
            if SCROLL != "PHerc0139" and os.path.exists(f"{d}/ink1um_ds8.jpg"):
                # second published ink layer (1.129 um scan): a per-winding overlay texture plus a close-up
                os.makedirs(f"{OUT}/ink1", exist_ok=True)
                i1 = Image.open(f"{d}/ink1um_ds8.jpg").convert("L")
                h1 = i1.copy()
                h1.thumbnail((HI_MAX, HI_MAX), Image.LANCZOS)
                h1.save(f"{OUT}/hi/{w}_1um.jpg", quality=84, optimize=True, progressive=True)
                i1.thumbnail((1200, 1200), Image.LANCZOS)
                i1.save(f"{OUT}/ink1/{w}.jpg", quality=80, optimize=True, progressive=True)
        ct = fetch_ct(w, src, (H, W))
        valid = ct > 0
        lo, hi_ = np.percentile(ct[valid], [1.5, 99.5]) if valid.any() else (0, 1)
        ctn = np.clip((ct - lo) / max(hi_ - lo, 1e-6), 0, 1)
        ctn[~valid] = 0
        ct_tiles[w] = Image.fromarray((ctn * 255).astype(np.uint8)).resize((tw, th), Image.LANCZOS)
        print(w, "tile", tw, th, "ct", ct.shape)
    pos, used = shelf_pack(sizes, ATLAS)
    if used > ATLAS:
        sys.exit(f"atlas overflow: {used} px > {ATLAS}; lower PPMM")
    atlas_h = 1 << int(np.ceil(np.log2(used)))
    atlases = {"ct": Image.new("L", (ATLAS, atlas_h), 0)}
    if not NO_INK:
        atlases["ink"] = Image.new("L", (ATLAS, atlas_h), 0)
    for w in names:
        for k in atlases:
            atlases[k].paste((ink_tiles if k == "ink" else ct_tiles)[w], pos[w])
    for k, im in atlases.items():
        im.save(f"{OUT}/atlas_{k}.jpg", quality=88, optimize=True, progressive=True)
        im.save(f"{OUT}/atlas_{k}.webp", quality=82, method=6)
    print("atlas", ATLAS, atlas_h, "used rows", used)

    # meshes
    blobs, wins, offset = [], [], 0
    for w, (x, y, z, v), src, meta in zip(names, grids, srcs, metas):
        H, W = v.shape
        rows = np.arange(0, H, STEP)
        cols = np.arange(0, W, STEP)
        if rows[-1] != H - 1:
            rows = np.append(rows, H - 1)
        if cols[-1] != W - 1:
            cols = np.append(cols, W - 1)
        xs, ys, zs, vs = (a[np.ix_(rows, cols)] for a in (x, y, z, v))
        # center on the fitted axis; x(z), y(z) keep the tilt in the data
        px = np.where(vs, xs - cx0, 0)
        py = np.where(vs, ys - cy0, 0)
        pz = np.where(vs, zs - zc, 0)
        st = np.stack([px, py, pz], -1).round()
        assert np.abs(st).max() < 32767, "coordinates overflow int16: use a coarser mesh"
        arr = st.astype(np.int16)
        arr[~vs] = -32768
        b = arr.tobytes()
        blobs.append(b)
        r = np.hypot(xs[vs] - np.polyval(ax, zs[vs]), ys[vs] - np.polyval(ay, zs[vs])) * VOX_MM
        x0, y0 = pos[w]
        tw, th = sizes[w]
        wins.append({
            "id": w,
            "seg": src["seg"],
            "rows": int(len(rows)), "cols": int(len(cols)),
            "rowIdx": rows.tolist(), "colIdx": cols.tolist(),
            "gridH": int(H), "gridW": int(W),
            "offset": offset, "bytes": len(b),
            "widthMM": round(W * CELL_MM, 2), "heightMM": round(H * CELL_MM, 2),
            "areaCM2": round((meta["area_vx2"] * VOX_MM * VOX_MM if "area_vx2" in meta else v.sum() * CELL_MM * CELL_MM) / 100, 2),
            "radiusMM": round(float(np.median(r)), 2),
            "atlas": [x0 / ATLAS, y0 / atlas_h, (x0 + tw) / ATLAS, (y0 + th) / atlas_h],
            **({} if NO_INK else {"hi": f"hi/{w}.jpg", "ink_source": src.get("ink_ds") or [f for f, _ in src["ink"] if "2.399um" in f][0]}),
            "mesh_source": src["mesh"],
        })
        offset += len(b)
    open(f"{OUT}/scroll.bin", "wb").write(b"".join(blobs))
    manifest = {
        "scroll": SCROLL,
        "voxelMM": VOX_MM, "cellMM": CELL_MM,
        "axis": {"centre": [cx0, cy0, zc], "dxdz": ax[0], "dydz": ay[0]},  # key as in the published data/*.json
        "atlas": {"w": ATLAS, "h": atlas_h, "ppmm": PPMM},
        "windings": wins,
        "totalWidthMM": round(sum(wi["widthMM"] for wi in wins), 1),
        "totalAreaCM2": round(sum(wi["areaCM2"] for wi in wins), 1),
        "vertices": int(offset // 6),
        "step": STEP,
    }
    json.dump(manifest, open(f"{OUT}/scroll.json", "w"), separators=(",", ":"))
    print("vertices", manifest["vertices"], "bin MB", offset / 1e6, "strip m", manifest["totalWidthMM"] / 1000,
          "area cm2", manifest["totalAreaCM2"])


def coarsen(folder, k, name="scroll_lite"):
    """Keep every k-th sampled row and column of each winding (always the first and last) of a built mesh."""
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from pack_bin import pack
    man = json.load(open(f"{folder}/scroll.json"))
    raw = np.fromfile(f"{folder}/scroll.bin", "<i2")
    blobs, offset, wins = [], 0, []
    for w in man["windings"]:
        R, C = w["rows"], w["cols"]
        a = raw[w["offset"] // 2:(w["offset"] + w["bytes"]) // 2].reshape(R, C, 3)
        ri = sorted(set(range(0, R, k)) | {R - 1})
        ci = sorted(set(range(0, C, k)) | {C - 1})
        b = np.ascontiguousarray(a[np.ix_(ri, ci)]).astype("<i2").tobytes()
        blobs.append(b)
        wins.append({**w, "rows": len(ri), "cols": len(ci), "rowIdx": [w["rowIdx"][i] for i in ri],
                     "colIdx": [w["colIdx"][i] for i in ci], "offset": offset, "bytes": len(b)})
        offset += len(b)
    data = b"".join(blobs)
    out = {**man, "windings": wins, "vertices": int(offset // 6), "step": man.get("step", 3) * k}
    open(f"{folder}/{name}.bin", "wb").write(data)
    open(f"{folder}/{name}.i16p.gz", "wb").write(pack(data))
    json.dump(out, open(f"{folder}/{name}.json", "w"), separators=(",", ":"))
    print("vertices", man["vertices"], "->", out["vertices"], "bin bytes", len(raw) * 2, "->", len(data),
          "packed", os.path.getsize(f"{folder}/{name}.i16p.gz"))


if __name__ == "__main__":
    if len(sys.argv) == 4 and sys.argv[1] == "--coarsen":
        coarsen(sys.argv[3].rstrip("/"), int(sys.argv[2]))
    else:
        main()