"""H&E patches of the 9,484 HEST training-pool spots (cache/hest/meta.csv), so that image encoders other than the
precomputed UNI2 / GigaPath features can be run on them (e.g. MedVTok).

Run:  HF_TOKEN=... .venv/bin/python -I 49_hest_patches.py [--workers 16] [--only <sample> ...]
Source: MahmoodLab/hest wsis/<sample>.tif (pyramidal tiled TIFF), level 0, read with HTTP range requests on the
resolved CDN URL (RangeFile) - only the tiles under the needed spots. (The patches/<sample>.h5 files cannot be read
selectively: their barcode chunks are interleaved with the image chunks, i.e. finding the rows reads the whole file.)
Patch = 224 x 224 level-0 px centred on the spot centre (obsm['spatial'] = full-res (x, y), cache/hest/spot_coords.csv),
no resizing: this is exactly the image behind the cached UNI2 features (re-encoding these crops with UNI2-h gives
cos 1.000 to cache/hest/he_uni2.npy, checked on TENX152 / MEND139; 0.5-um/px resampled crops give only 0.66-0.81).
Pixels outside the slide are white. Resumable per sample.
Output cache/hest/patches/<sample>.npz (img, idx) and, once all are present, cache/hest/he_img.npy (9484, 224, 224, 3)
uint8 in meta.csv row order (+ cache/hest/he_img_missing.csv, spots whose slide could not be read).
"""
import io
import os
import sys
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import pandas as pd
import requests
import tifffile

ROOT = os.path.dirname(os.path.abspath(__file__))
H = os.path.join(ROOT, "cache", "hest")
OUT = os.path.join(H, "patches")
os.makedirs(OUT, exist_ok=True)
WORKERS = int(sys.argv[sys.argv.index("--workers") + 1]) if "--workers" in sys.argv else 16
ONLY = sys.argv[sys.argv.index("--only") + 1:] if "--only" in sys.argv else None
M = pd.read_csv(os.path.join(H, "meta.csv"))
C = pd.read_csv(os.path.join(H, "spot_coords.csv"))
assert (C.idx.values == M.idx.values).all()


class RangeFile(io.RawIOBase):
    """Read-only, seekable file over HTTP range requests (block cache), for h5py. huggingface_hub's HfFileSystem hung
    in this environment (2026-10-28), plain range requests on the resolved CDN URL are fast."""
    BLOCK = 256 * 1024

    def __init__(self, url):
        hdr = {"Authorization": f"Bearer {os.environ['HF_TOKEN']}"}
        r = requests.get(url, headers={**hdr, "Range": "bytes=0-0"}, timeout=60)
        r.raise_for_status()
        self.url = r.url  # signed CDN URL after redirects
        self.size = int(r.headers["Content-Range"].split("/")[1])
        self.pos, self.cache, self.s = 0, {}, requests.Session()

    def _block(self, i):
        if i not in self.cache:
            a = i * self.BLOCK
            for t in range(5):
                try:
                    r = self.s.get(self.url, headers={"Range": f"bytes={a}-{min(self.size, a + self.BLOCK) - 1}"}, timeout=60)
                    r.raise_for_status()
                    break
                except requests.RequestException:
                    if t == 4:
                        raise
            self.cache[i] = r.content
        return self.cache[i]

    def readable(self):
        return True

    def seekable(self):
        return True

    def tell(self):
        return self.pos

    def seek(self, off, whence=0):
        self.pos = off if whence == 0 else self.pos + off if whence == 1 else self.size + off
        return self.pos

    def readinto(self, b):
        n = min(len(b), max(0, self.size - self.pos))
        out, p = bytearray(), self.pos
        while len(out) < n:
            blk = self._block(p // self.BLOCK)
            o = p % self.BLOCK
            take = blk[o:o + n - len(out)]
            out += take; p += len(take)
        b[:n] = out
        self.pos = p
        return n


def fetch(sid):
    out = os.path.join(OUT, f"{sid}.npz")
    if os.path.exists(out):
        return sid, "cached"
    c = C[C["sample"] == sid]
    f = RangeFile(f"https://huggingface.co/datasets/MahmoodLab/hest/resolve/main/wsis/{sid}.tif")
    with tifffile.TiffFile(f) as tf:
        pg = tf.pages[0]
        Hh, W = pg.shape[:2]
        th, tw = pg.tilelength, pg.tilewidth
        ntx = -(-W // tw)
        S = 224
        tiles = {}

        def tile(ty, tx):
            if (ty, tx) not in tiles:
                i = ty * ntx + tx
                f.seek(pg.dataoffsets[i])
                a = pg.decode(f.read(pg.databytecounts[i]), i, jpegtables=pg.jpegtables)[0]
                tiles[(ty, tx)] = a.reshape(th, tw, -1)[..., :3]
            return tiles[(ty, tx)]

        imgs = []
        for x, y in zip(c.x.values, c.y.values):
            x0, y0 = int(round(x - S / 2)), int(round(y - S / 2))
            reg = np.full((S, S, 3), 255, np.uint8)
            for ty in range(max(0, y0 // th), min(-(-Hh // th), (y0 + S - 1) // th + 1)):
                for tx in range(max(0, x0 // tw), min(ntx, (x0 + S - 1) // tw + 1)):
                    a = tile(ty, tx)
                    ys, xs = max(y0, ty * th), max(x0, tx * tw)
                    ye, xe = min(y0 + S, (ty + 1) * th, Hh), min(x0 + S, (tx + 1) * tw, W)
                    reg[ys - y0:ye - y0, xs - x0:xe - x0] = a[ys - ty * th:ye - ty * th, xs - tx * tw:xe - tx * tw]
            imgs.append(reg)
    np.savez(out + ".tmp.npz", img=np.stack(imgs), idx=c.idx.values)
    os.replace(out + ".tmp.npz", out)
    return sid, f"{len(imgs)} spots, crop {S}px, {len(tiles)} tiles, {sum(len(v) for v in f.cache.values()) / 1e6:.0f} MB"


samples = ONLY or M["sample"].unique().tolist()
with ThreadPoolExecutor(WORKERS) as ex:
    futs = {s: ex.submit(fetch, s) for s in samples}
    for n, (sid, fu) in enumerate(futs.items()):
        try:
            print(f"[{n + 1}/{len(samples)}] {sid}: {fu.result()[1]}", flush=True)
        except Exception as e:  # noqa: BLE001 - reported, spots become 'missing'
            print(f"[{n + 1}/{len(samples)}] {sid}: FAILED {type(e).__name__}: {e}", flush=True)
if ONLY:
    sys.exit()

img = np.zeros((len(M), 224, 224, 3), np.uint8)
have = np.zeros(len(M), bool)
for sid in M["sample"].unique():
    p = os.path.join(OUT, f"{sid}.npz")
    if os.path.exists(p):
        z = np.load(p)
        img[z["idx"]] = z["img"]; have[z["idx"]] = True
np.save(os.path.join(H, "he_img.npy"), img)
pd.DataFrame({"idx": np.flatnonzero(~have)}).to_csv(os.path.join(H, "he_img_missing.csv"), index=False)
print(f"saved he_img.npy {img.shape}; missing {(~have).sum()} spots")