"""HOT building tiles (OAM imagery CC BY 4.0, OSM labels ODbL): mosaic 2x2 adjacent z19 tiles (~0.3 m) -> 512 -> 256 (~0.6 m).
Outputs prep/hot_x.npy (M,256,256,3), prep/hot_y.npy (M,256,256) with 1 = building."""
import io, sys, os
from collections import defaultdict
import numpy as np, cv2, pyarrow.parquet as pq
from PIL import Image
from huggingface_hub import hf_hub_download
NSH = int(sys.argv[1]) if len(sys.argv) > 1 else 4
tiles = {}
for i in range(NSH):
    p = hf_hub_download("hotosm/vhr-building-segmentation", f"data/train-{i:05d}-of-00011.parquet", repo_type="dataset", local_dir="/workspace/hfd")
    t = pq.read_table(p, columns=["image", "mask", "tile_x", "tile_y", "project_id"]).to_pylist()
    for r in t:
        tiles[(r["project_id"], r["tile_x"], r["tile_y"])] = (r["image"]["bytes"], r["mask"]["bytes"])
    os.remove(p); print("shard", i, len(tiles), flush=True)
blocks = defaultdict(dict)
for (pid, x, y) in tiles: blocks[(pid, x // 2, y // 2)][(x % 2, y % 2)] = (pid, x, y)
X, Y = [], []
def dec(b, gray=False):
    im = Image.open(io.BytesIO(b)); a = np.asarray(im.convert("L" if gray else "RGB"))
    return a
for k, v in blocks.items():
    if len(v) < 4: continue
    try:
        img = np.zeros((512, 512, 3), np.uint8); msk = np.zeros((512, 512), np.uint8)
        for (dx, dy), key in v.items():
            ib, mb = tiles[key]; a = dec(ib); m = dec(mb, True)
            if a.shape[:2] != (256, 256) or m.shape != (256, 256): raise ValueError
            img[dy * 256:(dy + 1) * 256, dx * 256:(dx + 1) * 256] = a; msk[dy * 256:(dy + 1) * 256, dx * 256:(dx + 1) * 256] = (m > 0)
    except Exception:
        continue
    if msk.mean() < 0.01 or (img.max(2) == 0).mean() > 0.05: continue
    X.append(cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)); Y.append(cv2.resize(msk, (256, 256), interpolation=cv2.INTER_NEAREST))
X, Y = np.stack(X), np.stack(Y)
np.save("/workspace/prep/hot_x.npy", X); np.save("/workspace/prep/hot_y.npy", Y)
print("HOT_DONE", X.shape, "building frac", round(float(Y.mean()), 4), flush=True)