data/dataset.py

"""HoHo26k dataset.

Two surface paths share the per-tar WebDataset → `build_scene_input` pipeline:

  - `HoHoStreamingDataset` (IterableDataset): reads the raw HF tars, applies
    `build_scene_input` + `build_gt_verts` per sample. Used by the
    preprocessing CLI to write per-tar shards, and by inline eval when no
    cache exists.

  - `HoHoCachedDataset` (map-style Dataset): loads ALL semantic cache shards into memory
    at construction (small enough — 26k samples × ~tens of KB each fits in a
    few GB). Used for training and cached eval. Combine with
    `DistributedSampler` for DDP-safe, shuffle-every-epoch loading. Optional
    train-time augmentation transforms live in `data.transforms`.

Cache filename (one shard per input tar):
    `{cache_dir}/{split}.sem_v7.shard{NNN}-of-{TOTAL}.pkl`

The shard index `NNN` matches the input tar's `dataset-NNNN.tar` index; TOTAL
is the total number of tars for the split.
"""
from __future__ import annotations

import hashlib
import json
import pickle
import re
from pathlib import Path
from typing import Iterable, List, Optional

import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

CACHE_VERSION = "sem_v10"
SUPPORTED_CACHE_VERSIONS = {"sem_v7", "sem_v10", "sem_v11"}
# Cache versions for which GT wireframes have small connected components
# (chimneys: <=4 verts and <5 m^2 XY footprint) dropped before building verts_gt.
CHIMNEY_FILTER_VERSIONS = {"sem_v11"}
SHARD_FILENAME_TEMPLATE = "{split}.{ver}.shard{shard:03d}-of-{total:03d}.pkl"
SHARD_GLOB_TEMPLATE = "{split}.{ver}.shard*.pkl"
_SHARD_RE = re.compile(r"\.shard(\d{3})-of-(\d{3})\.pkl$")

from hoho2025.example_solutions import read_colmap_rec
from data.preprocess import build_scene_input, build_gt_verts, filter_small_components

# V11 chimney filter parameters: any connected component of the raw GT
# wireframe with <= V11_MAX_VERTS_PER_CC vertices AND XY footprint < V11_MAX_AREA_M2
# square metres is dropped (in raw world coords) before building verts_gt.
V11_MAX_VERTS_PER_CC = 4
V11_MAX_AREA_M2 = 2.0

DATASET_ID = "usm3d/hoho22k_2026_trainval"


# ---------------------------------------------------------------------------
# Raw WebDataset sample decoder (shared by both dataset classes)
# ---------------------------------------------------------------------------

def _parse_npy(raw_bytes):
    import io
    return np.load(io.BytesIO(raw_bytes))


def decode_wds_sample(sample: dict) -> dict:
    """Decode one raw WebDataset sample dict into the canonical HoHo record."""
    order_id_str = str(sample.get("__key__", "")).replace("order_", "")

    K_dict, R_dict, t_dict, pose_only_dict = {}, {}, {}, {}
    ade_dict, depth_dict, gestalt_dict = {}, {}, {}
    wf_vertices_list, wf_edges_list, wf_classifications_list = [], [], []
    colmap_zip_bytes = None

    for k, v in sample.items():
        k_lower = k.lower()
        if k_lower == "__key__":
            continue
        elif k_lower == "colmap.zip":
            colmap_zip_bytes = v
        elif k_lower.startswith("ade_") and k_lower.endswith(".png"):
            ade_dict[k_lower[4:-4]] = {"bytes": v}
        elif k_lower.startswith("depth_") and k_lower.endswith(".png"):
            depth_dict[k_lower[6:-4]] = {"bytes": v}
        elif k_lower.startswith("gestalt_") and k_lower.endswith(".png"):
            gestalt_dict[k_lower[8:-4]] = {"bytes": v}
        elif k_lower.endswith(".npy"):
            if k_lower.startswith("k_"):
                K_dict[k_lower[2:-4]] = _parse_npy(v)
            elif k_lower.startswith("r_"):
                R_dict[k_lower[2:-4]] = _parse_npy(v)
            elif k_lower.startswith("t_"):
                t_dict[k_lower[2:-4]] = _parse_npy(v)
            elif k_lower.startswith("pose_only_in_colmap_"):
                pose_only_dict[k_lower[20:-4]] = bool(_parse_npy(v))
            elif k_lower == "wf_vertices.npy":
                wf_vertices_list = _parse_npy(v).tolist()
            elif k_lower == "wf_edges.npy":
                wf_edges_list = _parse_npy(v).tolist()
            elif k_lower == "wf_classifications.npy":
                arr = _parse_npy(v)
                if arr.shape == ():
                    arr = arr.reshape((1,))
                wf_classifications_list = arr.tolist()

    img_ids_list = sorted(K_dict.keys())
    K_list, R_list, t_list, pose_only_list = [], [], [], []
    ade_list, depth_list, gestalt_list = [], [], []

    for img_id in img_ids_list:
        K_arr = K_dict.get(img_id, np.zeros((3, 3), dtype=np.float32))
        R_arr = R_dict.get(img_id, np.zeros((3, 3), dtype=np.float32))
        t_arr = np.array(t_dict.get(img_id, np.zeros((3, 1), dtype=np.float32)), dtype=np.float32)
        if t_arr.ndim == 1:
            t_arr = t_arr.reshape(3, 1)

        K_list.append(K_arr.tolist())
        R_list.append(R_arr.tolist())
        t_list.append(t_arr.tolist())
        pose_only_list.append(pose_only_dict.get(img_id, False))

        ade_list.append(ade_dict.get(img_id, None))
        depth_list.append(depth_dict.get(img_id, None))
        gestalt_list.append(gestalt_dict.get(img_id, None))

    return {
        "order_id": order_id_str,
        "image_ids": img_ids_list,
        "ade": ade_list,
        "depth": depth_list,
        "gestalt": gestalt_list,
        "K": K_list,
        "R": R_list,
        "t": t_list,
        "pose_only_in_colmap": pose_only_list,
        "wf_vertices": wf_vertices_list,
        "wf_edges": wf_edges_list,
        "wf_classifications": wf_classifications_list,
        "colmap": colmap_zip_bytes,
    }


# ---------------------------------------------------------------------------
# Tar discovery + shard naming
# ---------------------------------------------------------------------------

_TAR_RE = re.compile(r"data/([a-z]+)/dataset-(\d+)\.tar")


def _split_aliases(split: str) -> str:
    return {"val": "validation", "valid": "validation", "dev": "validation"}.get(split, split)


def discover_tars(cache_dir: Path, dataset_id: str, split: str) -> List[tuple[int, str]]:
    """Return [(tar_index, tar_local_path), …] sorted by tar_index."""
    split = _split_aliases(split)
    downloads_dir = Path(cache_dir) / "downloads"
    out: list[tuple[int, str]] = []
    for jf in downloads_dir.glob("*.json"):
        try:
            data = json.loads(jf.read_text())
        except json.JSONDecodeError:
            continue
        url = data.get("url", "")
        if not url.startswith(f"hf://datasets/{dataset_id}/data/{split}/"):
            continue
        m = _TAR_RE.search(url)
        if m is None:
            continue
        tar_idx = int(m.group(2))
        tar_path = str(jf)[:-5]  # strip .json
        out.append((tar_idx, tar_path))
    out.sort(key=lambda x: x[0])
    return out


def _normalize_cache_version(cache_version: str) -> str:
    ver = str(cache_version).strip().lower()
    if ver in {"v7", "7"}:
        ver = "sem_v7"
    elif ver in {"v10", "10"}:
        ver = "sem_v10"
    elif ver in {"v11", "11"}:
        ver = "sem_v11"
    if ver not in SUPPORTED_CACHE_VERSIONS:
        raise ValueError(
            f"Unsupported cache_version={cache_version!r}; "
            f"expected one of {sorted(SUPPORTED_CACHE_VERSIONS)}"
        )
    return ver


def shard_filename(split: str, shard_idx: int, total: int, cache_version: str = CACHE_VERSION) -> str:
    cache_version = _normalize_cache_version(cache_version)
    return SHARD_FILENAME_TEMPLATE.format(
        split=_split_aliases(split), ver=cache_version, shard=shard_idx, total=total
    )


def discover_shards(cache_dir: Path, split: str, cache_version: str = CACHE_VERSION) -> List[Path]:
    """Return existing semantic cache shard files for a split, sorted by shard index."""
    split = _split_aliases(split)
    cache_version = _normalize_cache_version(cache_version)
    paths = sorted(Path(cache_dir).glob(SHARD_GLOB_TEMPLATE.format(split=split, ver=cache_version)))
    paths.sort(key=lambda p: int(_SHARD_RE.search(p.name).group(1)))
    return paths


def _seed_from_order_id(order_id: str) -> int:
    """Stable 64-bit seed derived from order_id (independent of stream position)."""
    h = hashlib.sha1(order_id.encode("utf-8")).digest()
    return int.from_bytes(h[:8], "big", signed=False)


# ---------------------------------------------------------------------------
# Per-sample preprocessing
# ---------------------------------------------------------------------------

def process_sample(
    sample: dict,
    n_pts: int,
    k_verts: int,
    use_depth: bool,
    voxel_size_m: float = 0.0,
    cache_version: str = CACHE_VERSION,
) -> Optional[dict]:
    """Decode + featurise one raw WebDataset sample. Returns None on failure.

    For ``cache_version`` in :data:`CHIMNEY_FILTER_VERSIONS` (e.g. ``sem_v11``),
    small connected components of the raw GT wireframe are dropped before
    building ``verts_gt``. The original (unfiltered) ``wf_vertices_raw`` and
    ``wf_edges_raw`` are still saved alongside boolean keep masks so eval and
    visualisation can show the deleted parts.
    """
    order_id = sample.get("order_id", "")
    rng = np.random.default_rng(_seed_from_order_id(order_id))
    cache_version = _normalize_cache_version(cache_version)
    try:
        colmap_rec = read_colmap_rec(sample["colmap"])
        sample_dec = {**sample, "colmap": colmap_rec}

        scene = build_scene_input(
            sample_dec,
            n_pts=n_pts,
            use_depth=use_depth,
            voxel_size_m=voxel_size_m,
            rng=rng,
        )

        wf_v_raw = np.asarray(sample["wf_vertices"], dtype=np.float32)
        wf_e_raw = np.asarray(sample["wf_edges"], dtype=np.int64)
        if wf_v_raw.ndim == 1:
            wf_v_raw = wf_v_raw.reshape(0, 3)
        if wf_e_raw.ndim == 1:
            wf_e_raw = wf_e_raw.reshape(0, 2)

        if cache_version in CHIMNEY_FILTER_VERSIONS:
            keep_v, keep_e = filter_small_components(
                wf_v_raw, wf_e_raw,
                max_verts_per_cc=V11_MAX_VERTS_PER_CC,
                max_area_m2=V11_MAX_AREA_M2,
            )
        else:
            keep_v = np.ones(len(wf_v_raw), dtype=bool)
            keep_e = np.ones(len(wf_e_raw), dtype=bool)

        # Build verts_gt from kept vertices only. We do this by overriding the
        # `wf_vertices` field that build_gt_verts reads.
        kept_verts = wf_v_raw[keep_v]
        sample_for_gt = {**sample_dec, "wf_vertices": kept_verts.tolist()}
        gt = build_gt_verts(
            sample_for_gt,
            scene["bbox_center"],
            scene["bbox_scale"],
            k_verts,
        )

        record = {
            "scene_xyz":         torch.from_numpy(scene["scene_xyz"]),
            "scene_type_ids":    torch.from_numpy(scene["scene_type_ids"]),
            "scene_gestalt_ids": torch.from_numpy(scene["scene_gestalt_ids"]),
            "scene_gestalt_id2": torch.from_numpy(scene["scene_gestalt_id2"]),
            "scene_gestalt_w1":  torch.from_numpy(scene["scene_gestalt_w1"]),
            "scene_ade_ids":     torch.from_numpy(scene["scene_ade_ids"]),
            "scene_geom_conf":   torch.from_numpy(scene["scene_geom_conf"]),
            "scene_sem_conf":    torch.from_numpy(scene["scene_sem_conf"]),
            "scene_rgb":         torch.from_numpy(scene["scene_rgb"]),
            "verts_gt":          torch.from_numpy(gt),
            "bbox_center":       torch.from_numpy(scene["bbox_center"]),
            "bbox_scale":        torch.tensor(scene["bbox_scale"], dtype=torch.float32),
            "bbox_R":            torch.from_numpy(scene["bbox_R"]),
            "order_id":          order_id,
            "n_gt_verts":        int(keep_v.sum()),
            "wf_vertices_raw":   wf_v_raw,
            "wf_edges_raw":      wf_e_raw,
        }
        if cache_version in CHIMNEY_FILTER_VERSIONS:
            record["wf_vertex_kept_mask"] = keep_v
            record["wf_edge_kept_mask"] = keep_e
        return record
    except Exception as e:
        print(f"[dataset] skip sample {order_id}: {e}", flush=True)
        return None


# ---------------------------------------------------------------------------
# Streaming (preprocessing) dataset — one tar at a time
# ---------------------------------------------------------------------------

class HoHoStreamingDataset(IterableDataset):
    """Iterable view over raw HF tars; yields featurised samples.

    Used by the preprocessing CLI to build cache shards, and by HoHoCachedDataset
    transparently as a fallback if no cache exists.
    """

    def __init__(
        self,
        dataset_id: str = DATASET_ID,
        split: str = "train",
        cache_dir: Optional[str] = None,
        n_pts: int = 8192,
        k_verts: int = 64,
        use_depth: bool = True,
        voxel_size_m: float = 0.0,
        tar_paths: Optional[Iterable[str]] = None,
        n_samples: Optional[int] = None,
        cache_version: str = CACHE_VERSION,
    ):
        self.dataset_id = dataset_id
        self.split = _split_aliases(split)
        self.cache_dir = Path(cache_dir) if cache_dir else None
        self.n_pts = n_pts
        self.k_verts = k_verts
        self.use_depth = use_depth
        self.voxel_size_m = float(voxel_size_m)
        self.n_samples = n_samples
        self.cache_version = _normalize_cache_version(cache_version)

        if tar_paths is not None:
            self.tar_paths = list(tar_paths)
        else:
            if self.cache_dir is None:
                raise ValueError("HoHoStreamingDataset needs cache_dir or explicit tar_paths.")
            tars = discover_tars(self.cache_dir, dataset_id, self.split)
            self.tar_paths = [path for _, path in tars]
            if not self.tar_paths:
                raise ValueError(
                    f"No cached tar files found for split {self.split} in "
                    f"{self.cache_dir}/downloads. Download the dataset first."
                )

    def __iter__(self):
        import webdataset as wds

        ds = wds.WebDataset(self.tar_paths, shardshuffle=False, empty_check=False)
        ds = ds.map(decode_wds_sample)

        emitted = 0
        for raw in ds:
            if self.n_samples is not None and emitted >= self.n_samples:
                break
            item = process_sample(
                raw, self.n_pts, self.k_verts, self.use_depth,
                voxel_size_m=self.voxel_size_m,
                cache_version=self.cache_version,
            )
            if item is not None:
                yield item
                emitted += 1


# ---------------------------------------------------------------------------
# Cached (training/eval) dataset — map-style, all shards in memory
# ---------------------------------------------------------------------------

class HoHoCachedDataset(Dataset):
    """Map-style dataset that loads all semantic cache shards into memory.

    Use with `torch.utils.data.distributed.DistributedSampler(shuffle=True)`
    to get unique-per-rank, reshuffled-every-epoch ordering across DDP ranks.

    Args:
        cache_dir: directory containing `{split}.{cache_version}.shard*.pkl` files.
        split: split name (`train`, `validation`).
        n_pts: model-input point count (= transformer query count).
        augment: if True, apply training augmentations on each `__getitem__`
            (random subsample, flip, yaw rotation, jitter). If False, take a
            deterministic prefix slice (eval/inference).
        flip / yaw / jitter_sigma: enable/strength of individual augmentations.
        n_samples: optional cap for debug runs.
    """

    def __init__(
        self,
        cache_dir: str,
        split: str = "train",
        n_pts: int = 8192,
        augment: bool = False,
        flip: bool = True,
        yaw: bool = True,
        jitter_sigma: float = 0.01,
        n_samples: Optional[int] = None,
        cache_version: str = CACHE_VERSION,
    ):
        super().__init__()
        self.cache_dir = Path(cache_dir)
        self.split = _split_aliases(split)
        self.cache_version = _normalize_cache_version(cache_version)
        self.n_pts = int(n_pts)
        self.augment = bool(augment)
        self.flip = bool(flip)
        self.yaw = bool(yaw)
        self.jitter_sigma = float(jitter_sigma)
        self.n_samples = n_samples

        shard_paths = discover_shards(self.cache_dir, self.split, self.cache_version)
        if not shard_paths:
            raise FileNotFoundError(
                f"No {self.cache_version} cache shards found at {self.cache_dir} for split {self.split}.\n"
                f"Run preprocessing first, e.g.:\n"
                f"    python -m data.dataset --split {self.split} --cache_dir {self.cache_dir} "
                f"--cache_version {self.cache_version} --shard_index 0 --total_shards 1"
            )

        self.samples: list[dict] = []
        for p in shard_paths:
            with open(p, "rb") as f:
                self.samples.extend(pickle.load(f))
        if self.n_samples is not None:
            self.samples = self.samples[: self.n_samples]

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> dict:
        # Avoid mutating the in-memory cache: shallow-copy the dict and the
        # per-point tensors that we will swap in. Other tensors are read-only.
        from data.transforms import apply_train_transforms, apply_eval_transforms

        src = self.samples[idx]
        out = {k: v for k, v in src.items()}
        # `verts_gt` may be rotated/flipped; clone to avoid aliasing the cache copy.
        if "verts_gt" in out and torch.is_tensor(out["verts_gt"]):
            out["verts_gt"] = out["verts_gt"].clone()

        if self.augment:
            return apply_train_transforms(
                out,
                n_pts=self.n_pts,
                yaw=self.yaw,
                flip=self.flip,
                jitter_sigma=self.jitter_sigma,
            )
        return apply_eval_transforms(out, n_pts=self.n_pts)


# ---------------------------------------------------------------------------
# Cache builder (run via SLURM array; one shard per input tar)
# ---------------------------------------------------------------------------

def build_shards(
    cache_dir: str,
    split: str,
    shard_index: int,
    total_shards: int,
    dataset_id: str = DATASET_ID,
    n_pts: int = 8192,
    k_verts: int = 64,
    use_depth: bool = True,
    voxel_size_m: float = 0.0,
    cache_version: str = CACHE_VERSION,
) -> List[Path]:
    """Build the shards assigned to this array task.

    Tar selection is `tars[shard_index::total_shards]` over the sorted tar
    list. Each tar produces exactly one output shard file whose `NNN` matches
    the input `dataset-NNNN.tar` index.
    """
    cache_dir_p = Path(cache_dir)
    cache_dir_p.mkdir(parents=True, exist_ok=True)
    cache_version = _normalize_cache_version(cache_version)

    tars = discover_tars(cache_dir_p, dataset_id, split)
    if not tars:
        raise ValueError(
            f"No cached tar files for split {split} in {cache_dir_p}/downloads."
        )

    # `total_tars` is used for the {NNN}-of-{TOTAL} suffix in filenames; we use
    # the maximum tar index +1 so that shard names line up with the dataset's
    # native `dataset-NNNN.tar` numbering even if some tars are missing locally.
    total_tars = max(t for t, _ in tars) + 1
    my_tars = tars[shard_index::total_shards]
    print(
        f"[shard] split={split} task {shard_index}/{total_shards}: "
        f"{len(my_tars)} tars (of {len(tars)} total)",
        flush=True,
    )

    written: list[Path] = []
    for tar_idx, tar_path in my_tars:
        out_path = cache_dir_p / shard_filename(split, tar_idx, total_tars, cache_version)
        if out_path.exists():
            print(f"[shard] skip (exists) {out_path.name}", flush=True)
            written.append(out_path)
            continue

        ds = HoHoStreamingDataset(
            dataset_id=dataset_id,
            split=split,
            cache_dir=str(cache_dir_p),
            n_pts=n_pts,
            k_verts=k_verts,
            use_depth=use_depth,
            voxel_size_m=voxel_size_m,
            tar_paths=[tar_path],
            cache_version=cache_version,
        )
        rows = list(ds)
        with open(out_path, "wb") as f:
            pickle.dump(rows, f)
        count_path = out_path.with_suffix(out_path.suffix + ".count")
        count_path.write_text(str(len(rows)))
        written.append(out_path)
        print(f"[shard] wrote {out_path.name}  ({len(rows)} samples)", flush=True)

    return written


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import argparse
    from config_utils import load_config, resolve_arg, DEFAULT_CONFIG_PATH

    p = argparse.ArgumentParser(description="Build HoHo semantic cache shards (one per input tar)")
    p.add_argument("--config", default=DEFAULT_CONFIG_PATH)
    p.add_argument("--split", default="train")
    p.add_argument("--cache_dir", default=None)
    p.add_argument("--dataset_id", default=None)
    p.add_argument("--n_pts_cache", type=int, default=None,
                   help="Cache point count (default 8192 = 4096 COLMAP/camera + 4096 depth)")
    p.add_argument("--k_verts", type=int, default=None)
    p.add_argument("--use_depth", action="store_true", default=None)
    p.add_argument("--no_depth", dest="use_depth", action="store_false")
    p.add_argument("--shard_index", type=int, default=0,
                   help="This task's index in the SLURM array (0-based)")
    p.add_argument("--total_shards", type=int, default=1,
                   help="Total array tasks; tars are split round-robin across these")
    p.add_argument("--voxel_size_m", type=float, default=None,
                   help="Metric voxel size (m) for joint downsampling; 0 disables")
    p.add_argument("--cache_version", default=CACHE_VERSION, choices=sorted(SUPPORTED_CACHE_VERSIONS),
                   help="Semantic cache filename version to write")
    args = p.parse_args()

    config, _ = load_config(args.config)
    dataset_id = resolve_arg(args.dataset_id, config, "dataset.id", DATASET_ID)
    cache_dir = resolve_arg(args.cache_dir, config, "paths.cache_dir", "cache/")
    n_pts_cache = resolve_arg(args.n_pts_cache, config, "model.n_pts_cache", 8192)
    k_verts = resolve_arg(args.k_verts, config, "model.k_verts", 64)
    use_depth = resolve_arg(args.use_depth, config, "train.use_depth", True)
    voxel_size_m = resolve_arg(args.voxel_size_m, config, "model.voxel_size_m", 0.0)

    build_shards(
        cache_dir=cache_dir,
        split=args.split,
        shard_index=args.shard_index,
        total_shards=args.total_shards,
        dataset_id=dataset_id,
        n_pts=n_pts_cache,
        k_verts=k_verts,
        use_depth=use_depth,
        voxel_size_m=voxel_size_m,
        cache_version=args.cache_version,
    )


# ---------------------------------------------------------------------------
# DataLoader collate
# ---------------------------------------------------------------------------

def collate_fn(batch: list) -> dict:
    """Stack fixed-shape tensor fields; keep ragged / string fields as lists."""
    if batch and "scene_rgb" not in batch[0]:
        for b in batch:
            b["scene_rgb"] = torch.zeros((*b["scene_xyz"].shape[:-1], 3), dtype=torch.uint8)
    tensor_keys = (
        "scene_xyz", "scene_type_ids",
        "scene_gestalt_ids", "scene_gestalt_id2", "scene_gestalt_w1",
        "scene_ade_ids", "scene_geom_conf", "scene_sem_conf",
        "scene_rgb",
        "verts_gt", "bbox_center", "bbox_scale", "bbox_R",
    )
    out = {k: torch.stack([b[k] for b in batch]) for k in tensor_keys}
    out["order_id"] = [b["order_id"] for b in batch]
    out["n_gt_verts"] = [b["n_gt_verts"] for b in batch]
    out["wf_vertices_raw"] = [b.get("wf_vertices_raw") for b in batch]
    out["wf_edges_raw"] = [b.get("wf_edges_raw") for b in batch]
    out["wf_vertex_kept_mask"] = [b.get("wf_vertex_kept_mask") for b in batch]
    out["wf_edge_kept_mask"] = [b.get("wf_edge_kept_mask") for b in batch]
    return out