code_snapshots/WiSER_public_release_code/wiser/data/dataset.py

"""Per-triple datasets for WiSER CIR set prediction.

Two datasets live here:

1. `TriplePathDataset`: minimal direct-record wrapper for unit tests.
2. `MultiSceneTripleDataset`: real WiSER-backed dataset. Uses the EXACT path
   scheme WiSER uses (see resolve_* helpers below): scene meta at
   `wireless_root / scene_id / f"scene_{dataset_tag}_meta.pt"` and mesh
   voxel cache at `scene3d_root / scene_id / "voxel_10cm" / "mesh_voxel_cache_100mm.pt"`.

Every sample emits the checkpoint-compatible `csi_path_*` target schema
produced by `build_single_cell_target`. The legacy `merged_*` names live
only inside the per-cell kernel in `csi_path_targets.py` and must not appear
in a dataset sample dict.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from typing import Any, Mapping, Sequence

import h5py
import numpy as np
import torch
from torch.utils.data import Dataset

from .csi_path_targets import MergedPathTargetConfig, build_single_cell_target, merge_paths_by_delay_bin

log = logging.getLogger(__name__)


# -----------------------------------------------------------------------------
# Synthetic path (TriplePathDataset)
# -----------------------------------------------------------------------------


@dataclass(slots=True)
class TripleRecord:
    scene_id: str
    tx_id: str
    rx_id: str
    tx_xyz_norm: np.ndarray
    rx_xyz_norm: np.ndarray
    scene_extent: np.ndarray
    raw_delay_ns: np.ndarray
    raw_power_linear: np.ndarray
    voxel_level: Any | None = None


@dataclass(slots=True)
class TriplePathDataset(Dataset):
    records: Sequence[TripleRecord]
    target_config: MergedPathTargetConfig = field(default_factory=MergedPathTargetConfig)

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

    def __getitem__(self, idx: int) -> dict[str, Any]:
        rec = self.records[idx]
        targets = build_single_cell_target(
            delay_ns=rec.raw_delay_ns,
            power_linear=rec.raw_power_linear,
            config=self.target_config,
        )
        return {
            "scene_id": rec.scene_id,
            "tx_id": rec.tx_id,
            "rx_id": rec.rx_id,
            "tx_xyz_norm": rec.tx_xyz_norm,
            "rx_xyz_norm": rec.rx_xyz_norm,
            "scene_extent": rec.scene_extent,
            "voxel_level": rec.voxel_level,
            **targets,
        }


def records_from_manifest(manifest_rows: Sequence[Mapping[str, Any]]) -> list[TripleRecord]:
    return [
        TripleRecord(
            scene_id=str(row["scene_id"]),
            tx_id=str(row["tx_id"]),
            rx_id=str(row["rx_id"]),
            tx_xyz_norm=np.asarray(row["tx_xyz_norm"], dtype=np.float32),
            rx_xyz_norm=np.asarray(row["rx_xyz_norm"], dtype=np.float32),
            scene_extent=np.asarray(row.get("scene_extent", [1.0, 1.0, 1.0]), dtype=np.float32),
            raw_delay_ns=np.asarray(row.get("raw_delay_ns", []), dtype=np.float64),
            raw_power_linear=np.asarray(row.get("raw_power_linear", []), dtype=np.float64),
            voxel_level=row.get("voxel_level"),
        )
        for row in manifest_rows
    ]


# -----------------------------------------------------------------------------
# Real WiSER dataset path (MultiSceneTripleDataset)
# -----------------------------------------------------------------------------


DEFAULT_WIRELESS_ROOT = Path(os.environ.get("WISER_WIRELESS_ROOT", "data/wireless"))
DEFAULT_SCENE3D_ROOT = Path(os.environ.get("WISER_SCENE3D_ROOT", "data/scene3d"))
DEFAULT_DATASET_TAG = "voxel_original_csi_path_10cm_1e6"


def resolve_scene_meta_path(scene_id: str, wireless_root: Path, dataset_tag: str) -> Path:
    """WiSER convention: wireless_root / scene_id / scene_{tag}_meta.pt."""
    return Path(wireless_root) / scene_id / f"scene_{dataset_tag}_meta.pt"


def resolve_voxel_cache_path(scene_id: str, scene3d_root: Path) -> Path:
    """WiSER convention: scene3d_root / scene_id / voxel_10cm / mesh_voxel_cache_100mm.pt."""
    return Path(scene3d_root) / scene_id / "voxel_10cm" / "mesh_voxel_cache_100mm.pt"


def _normalize_xyz(xyz: np.ndarray, coord_min: np.ndarray, coord_max: np.ndarray) -> np.ndarray:
    denom = np.maximum(coord_max - coord_min, 1e-6)
    return (2.0 * (xyz - coord_min) / denom - 1.0).astype(np.float32)


def _read_txt_positions(path: Path) -> np.ndarray:
    """Read the WiSER 'tx_positions' text file (3 floats per line, whitespace-separated).

    Lines starting with '#' are skipped.
    """
    rows: list[list[float]] = []
    with open(path, "r") as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split()
            if len(parts) >= 3:
                rows.append([float(parts[0]), float(parts[1]), float(parts[2])])
    if not rows:
        return np.zeros((0, 3), dtype=np.float32)
    return np.asarray(rows, dtype=np.float32)


@lru_cache(maxsize=32)
def _load_scene_meta_cached(path_str: str) -> dict[str, np.ndarray] | None:
    p = Path(path_str)
    if not p.exists():
        return None
    meta = torch.load(p, map_location="cpu", weights_only=False)

    def _arr(v) -> np.ndarray:
        if torch.is_tensor(v):
            return v.detach().cpu().numpy()
        return np.asarray(v)

    out: dict[str, np.ndarray] = {}
    out["coord_min"] = _arr(meta["coord_min"]).astype(np.float32)
    out["coord_max"] = _arr(meta["coord_max"]).astype(np.float32)

    # Resolve tx_positions: prefer in-dict tensor, else read from tx_positions_path.
    tx_pos = meta.get("tx_positions")
    if tx_pos is not None:
        out["tx_positions"] = _arr(tx_pos).astype(np.float32)
    else:
        tx_path = meta.get("tx_positions_path")
        if tx_path and Path(tx_path).exists():
            out["tx_positions"] = _read_txt_positions(Path(tx_path))
        else:
            out["tx_positions"] = np.zeros((0, 3), dtype=np.float32)

    # Resolve rx_xyz: either in-dict or computed from nx/ny/z_values/cell_size.
    rx_xyz = meta.get("rx_xyz")
    if rx_xyz is not None:
        out["rx_xyz"] = _arr(rx_xyz).astype(np.float32)
    else:
        nx = int(meta.get("nx", 0))
        ny = int(meta.get("ny", 0))
        cell = float(meta.get("cell_size", 0.0))
        z_vals = meta.get("z_values")
        if nx > 0 and ny > 0 and cell > 0.0 and z_vals is not None:
            z_arr = _arr(z_vals).astype(np.float32)
            num_z = int(z_arr.shape[0])
            x0 = float(meta.get("x_min", out["coord_min"][0]))
            y0 = float(meta.get("y_min", out["coord_min"][1]))
            margin = float(meta.get("xy_margin", 0.0))
            xs = x0 + margin + cell * (np.arange(nx, dtype=np.float32) + 0.5)
            ys = y0 + margin + cell * (np.arange(ny, dtype=np.float32) + 0.5)
            gx, gy = np.meshgrid(xs, ys, indexing="xy")  # [ny, nx]
            rx = np.zeros((num_z, ny, nx, 3), dtype=np.float32)
            rx[..., 0] = gx
            rx[..., 1] = gy
            for zi, zv in enumerate(z_arr):
                rx[zi, ..., 2] = zv
            out["rx_xyz"] = rx
        else:
            out["rx_xyz"] = np.zeros((0, 0, 0, 3), dtype=np.float32)

    for k in ("nx", "ny", "cell_size"):
        if k in meta:
            out[k] = meta[k]
    return out


@lru_cache(maxsize=32)
def _load_voxel_cache_cached(path_str: str, voxel_channels: int) -> dict[str, torch.Tensor] | None:
    """Load an WiSER mesh voxel cache and derive `{"feats", "coords"}`.

    Real WiSER caches expose `coords, counts, centers_world, colors, origin,
    voxel_size_m` rather than a precomputed `feats` tensor. We concatenate
    counts/centers_world/colors into a 7-D per-voxel feature vector (matching
    WiSER's `voxel_feature_dim=16` convention, padded to the config channel
    count). When a direct `feats` tensor IS present (some preprocessed dumps)
    we use it as-is.
    """
    p = Path(path_str)
    if not p.exists():
        return None
    obj = torch.load(p, map_location="cpu", weights_only=False)

    feats_key = next((k for k in ("feats", "voxel_feats", "features") if k in obj), None)
    coords_key = next((k for k in ("coords", "voxel_coords", "coordinates") if k in obj), None)
    if feats_key is not None:
        feats = obj[feats_key]
        coords = obj[coords_key] if coords_key else obj["coords"]
    else:
        # Derive features from the WiSER cache layout.
        coords = obj.get("coords")
        centers = obj.get("centers_world")
        counts = obj.get("counts")
        colors = obj.get("colors")
        if coords is None or centers is None:
            return None
        parts = [centers.float()]
        if counts is not None:
            parts.append(counts.float().unsqueeze(-1))
        if colors is not None:
            parts.append(colors.float())
        feats = torch.cat(parts, dim=-1)

    feats = feats.float()
    coords = coords.float() if torch.is_tensor(coords) else torch.as_tensor(coords, dtype=torch.float32)
    if feats.dim() == 2:
        feats = feats.unsqueeze(0)
    if coords.dim() == 2:
        coords = coords.unsqueeze(0)

    # Pad/truncate channel dim to `voxel_channels` so the backbone forward fits.
    if feats.shape[-1] < voxel_channels:
        pad = voxel_channels - feats.shape[-1]
        feats = torch.cat([feats, torch.zeros((*feats.shape[:-1], pad), dtype=feats.dtype)], dim=-1)
    elif feats.shape[-1] > voxel_channels:
        feats = feats[..., :voxel_channels]
    return {"feats": feats, "coords": coords}


@lru_cache(maxsize=256)
def _open_h5_group(tx_path: str, z_key: str) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None:
    if not Path(tx_path).exists():
        return None
    with h5py.File(str(tx_path), "r") as h:
        if z_key not in h:
            return None
        grp = h[z_key]
        return (
            np.asarray(grp["cell_ptr"], dtype=np.int64),
            np.asarray(grp["cell_path_count"], dtype=np.int64),
            np.asarray(grp["path_delay_ns"], dtype=np.float64),
            np.asarray(grp["path_power_linear"], dtype=np.float64),
        )


class AssetMissingError(RuntimeError):
    """Raised when a required real WiSER asset is not reachable and
    `allow_synthetic_fallback` is False."""


class MultiSceneTripleDataset(Dataset):
    """Real WiSER-backed triple dataset.

    Loaders follow the WiSER release path scheme:
        scene meta   = wireless_root / scene_id / scene_{dataset_tag}_meta.pt
        voxel cache  = scene3d_root / scene_id / voxel_10cm / mesh_voxel_cache_100mm.pt

    `allow_synthetic_fallback=False` (default) causes missing assets to
    raise `AssetMissingError`, naming the exact missing path. This is the
    mode required by plan validation. Setting it to True restores the
    Round-2 silent fallback for local smoke work only and should never be
    used for any report whose metrics are meant to support an AC gate.

    Manifest-wide zero-path and truncation statistics are precomputed from
    the HDF5 supervision during `__init__` and exposed via `manifest_stats()`.
    This is genuine dataset-level evidence, not a run-time access counter.
    """

    def __init__(
        self,
        manifest: Mapping[str, Any],
        *,
        target_config: MergedPathTargetConfig | None = None,
        voxel_channels: int = 256,
        wireless_root: Path | str | None = None,
        scene3d_root: Path | str | None = None,
        dataset_tag: str = DEFAULT_DATASET_TAG,
        allow_synthetic_fallback: bool = False,
        precompute_stats: bool = True,
    ) -> None:
        self.manifest = manifest
        self.triples = list(manifest.get("triples", []))
        self.target_config = target_config or MergedPathTargetConfig()
        self.voxel_channels = int(voxel_channels)
        self.wireless_root = Path(wireless_root) if wireless_root is not None else DEFAULT_WIRELESS_ROOT
        self.scene3d_root = Path(scene3d_root) if scene3d_root is not None else DEFAULT_SCENE3D_ROOT
        self.dataset_tag = str(dataset_tag)
        self.allow_synthetic_fallback = bool(allow_synthetic_fallback)
        self._stats = None
        self._fallbacks = {"scene_meta": 0, "voxel_cache": 0}
        if precompute_stats:
            self._stats = self._scan_manifest_stats()

    # ------------- asset resolution -------------

    def _meta_path(self, scene_id: str) -> Path:
        return resolve_scene_meta_path(scene_id, self.wireless_root, self.dataset_tag)

    def _voxel_path(self, scene_id: str) -> Path:
        return resolve_voxel_cache_path(scene_id, self.scene3d_root)

    def _get_scene_meta(self, scene_id: str) -> dict[str, np.ndarray]:
        meta = _load_scene_meta_cached(str(self._meta_path(scene_id)))
        if meta is not None:
            return meta
        if not self.allow_synthetic_fallback:
            raise AssetMissingError(
                f"scene_meta missing for {scene_id} at {self._meta_path(scene_id)}. "
                f"Set allow_synthetic_fallback=True for local smoke work."
            )
        self._fallbacks["scene_meta"] += 1
        return {
            "coord_min": np.zeros((3,), dtype=np.float32),
            "coord_max": np.ones((3,), dtype=np.float32),
            "tx_positions": np.zeros((64, 3), dtype=np.float32),
            "rx_xyz": np.zeros((6, 48, 48, 3), dtype=np.float32),
        }

    def _get_voxel_cache(self, scene_id: str) -> dict[str, torch.Tensor]:
        vox = _load_voxel_cache_cached(str(self._voxel_path(scene_id)), self.voxel_channels)
        if vox is not None:
            return vox
        if not self.allow_synthetic_fallback:
            raise AssetMissingError(
                f"voxel_cache missing for {scene_id} at {self._voxel_path(scene_id)}. "
                f"Set allow_synthetic_fallback=True for local smoke work."
            )
        self._fallbacks["voxel_cache"] += 1
        rng = np.random.default_rng(abs(hash(scene_id)) % (2 ** 32))
        N = 128
        return {
            "feats": torch.from_numpy(rng.standard_normal((1, N, self.voxel_channels)).astype(np.float32)),
            "coords": torch.from_numpy(rng.random((1, N, 3)).astype(np.float32)),
        }

    # ------------- manifest-wide statistics (precomputed at init) -------------

    def _scan_manifest_stats(self) -> dict[str, int]:
        """Stream the full manifest once through HDF5 to compute real stats."""
        stats = {
            "scanned": 0,
            "zero_path": 0,
            "nonzero_path": 0,
            "truncated_samples": 0,
            "total_merged_paths": 0,
            "total_source_paths": 0,
            "missing_h5": 0,
        }
        for t in self.triples:
            stats["scanned"] += 1
            tx_path = t.get("tx_path")
            z_key = t.get("z_key")
            rx_flat = int(t["rx_flat_index"])
            if not tx_path or not z_key:
                stats["missing_h5"] += 1
                stats["zero_path"] += 1
                continue
            grp = _open_h5_group(str(tx_path), str(z_key))
            if grp is None:
                stats["missing_h5"] += 1
                stats["zero_path"] += 1
                continue
            cell_ptr, cell_path_count, all_delay, all_power = grp
            if rx_flat >= cell_path_count.shape[0]:
                stats["missing_h5"] += 1
                stats["zero_path"] += 1
                continue
            start = int(cell_ptr[rx_flat])
            count = int(cell_path_count[rx_flat])
            if count <= 0:
                stats["zero_path"] += 1
                continue
            per = merge_paths_by_delay_bin(
                delay_ns=all_delay[start:start + count],
                power_linear=all_power[start:start + count],
                config=self.target_config,
            )
            merged = int(per["merged_count"])
            if merged == 0:
                stats["zero_path"] += 1
            else:
                stats["nonzero_path"] += 1
            stats["total_merged_paths"] += merged
            stats["total_source_paths"] += int(per["source_path_count"])
            if int(per["truncated_count"]) > 0:
                stats["truncated_samples"] += 1
        return stats

    def manifest_stats(self) -> dict[str, Any]:
        N = max(int(self._stats["scanned"]) if self._stats else 0, 1)
        base = dict(self._stats or {})
        base["zero_path_rate"] = (self._stats["zero_path"] / N) if self._stats else 0.0
        base["truncated_rate"] = (self._stats["truncated_samples"] / N) if self._stats else 0.0
        base["scene_meta_fallbacks"] = self._fallbacks["scene_meta"]
        base["voxel_cache_fallbacks"] = self._fallbacks["voxel_cache"]
        return base

    # ------------- sample fetch -------------

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

    def __getitem__(self, idx: int) -> dict[str, Any]:
        # Round-2 (AC-11 speed): cache per-triple deterministic payload
        # (coords + merged-delay-bin targets) across epochs. Voxel payload
        # stays outside the cache because it is shared across (scene, TX, RX)
        # triples and already memoised by `_get_voxel_cache`.
        cached = getattr(self, "_triple_payload_cache", None)
        if cached is None:
            self._triple_payload_cache = cached = {}
        entry = cached.get(idx)
        if entry is None:
            t = self.triples[idx]
            scene_id = str(t["scene_id"])
            tx_idx = int(t["tx_idx"])
            z_idx = int(t["z_idx"])
            rx_flat = int(t["rx_flat_index"])

            meta = self._get_scene_meta(scene_id)
            coord_min = meta["coord_min"]
            coord_max = meta["coord_max"]
            scene_extent = np.maximum(coord_max - coord_min, 1e-6).astype(np.float32)

            tx_positions = meta["tx_positions"]
            tx_xyz_metric = tx_positions[tx_idx] if tx_idx < tx_positions.shape[0] else np.zeros((3,), dtype=np.float32)
            tx_xyz_norm = _normalize_xyz(tx_xyz_metric, coord_min, coord_max)

            rx_xyz_arr = meta["rx_xyz"]
            if rx_xyz_arr.ndim == 4 and z_idx < rx_xyz_arr.shape[0]:
                flat = rx_xyz_arr[z_idx].reshape(-1, 3)
                rx_xyz_metric = flat[rx_flat % flat.shape[0]]
            else:
                rx_xyz_metric = np.zeros((3,), dtype=np.float32)
            rx_xyz_norm = _normalize_xyz(rx_xyz_metric, coord_min, coord_max)

            tx_path = t.get("tx_path")
            z_key = t.get("z_key")
            raw_delay = np.empty((0,), dtype=np.float64)
            raw_power = np.empty((0,), dtype=np.float64)
            if tx_path and z_key:
                grp = _open_h5_group(str(tx_path), str(z_key))
                if grp is not None:
                    cell_ptr, cell_path_count, all_delay, all_power = grp
                    if rx_flat < cell_path_count.shape[0]:
                        start = int(cell_ptr[rx_flat])
                        count = int(cell_path_count[rx_flat])
                        raw_delay = all_delay[start:start + count]
                        raw_power = all_power[start:start + count]

            targets = build_single_cell_target(
                delay_ns=raw_delay, power_linear=raw_power, config=self.target_config,
            )
            entry = {
                "scene_id": scene_id,
                "tx_id": int(tx_idx),
                "rx_id": int(rx_flat),
                "tx_xyz_norm": tx_xyz_norm,
                "rx_xyz_norm": rx_xyz_norm,
                "scene_extent": scene_extent,
                **targets,
            }
            cached[idx] = entry
        out = dict(entry)
        out["voxel_level"] = self._get_voxel_cache(entry["scene_id"])
        return out


__all__ = [
    "AssetMissingError",
    "DEFAULT_DATASET_TAG",
    "DEFAULT_SCENE3D_ROOT",
    "DEFAULT_WIRELESS_ROOT",
    "MultiSceneTripleDataset",
    "TripleRecord",
    "TriplePathDataset",
    "records_from_manifest",
    "resolve_scene_meta_path",
    "resolve_voxel_cache_path",
]