Blame for src/data/multi_dataloader.py

"""
Multi-dataset data loader for multi-body dynamics.
Loads sequences from multiple dataset directories, with train/val/test
membership defined by a split file.
"""

import os
import json
import time
import concurrent.futures as futures
import h5py
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, Sampler, RandomSampler
from torch.utils.data.distributed import DistributedSampler
from typing import Dict, List, Tuple, Optional
import glob
from pathlib import Path
import random
from tqdm import tqdm
import sys
from functools import lru_cache
import logging
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.noise_utils import apply_random_walk_noise_to_batch, apply_independent_noise
from data.object_scale import ObjectScale, check_object_counts
from data.partial_view import PartialView, quat_to_rotmat
from data.scene_frames import frame_dirs as scene_frame_dirs, is_packed, load_frame
import zlib

logger = logging.getLogger(__name__)

# Padding value for positions: zero contributes nothing to
# pooling/aggregation operations.
PADDING_VALUE = 0.0


# Global LRU cache for frequently loaded numpy arrays
# (~50000 entries at ~5KB each, roughly 250MB).
@lru_cache(maxsize=50000)
def _cached_numpy_load(filepath: str) -> np.ndarray:
    """Cached numpy array loading. Uses filepath as cache key."""
    return np.load(filepath)


def cached_load_npy(path: Path) -> np.ndarray:
    """Load numpy array with LRU caching."""
    return _cached_numpy_load(str(path))


def check_scene_set(config: dict) -> str:
    """data.scene_set must name the scene set the split file declares ('benchmark' when either is absent).

    'complete' splits (data_generation/build_complete_splits.py) leave out the MOVi scenes whose metadata omits simulated
    bodies; 'benchmark' is the published scene list. Returns the scene set.
    """
    want = (config.get('data', {}) or {}).get('scene_set', 'benchmark')
    if want not in ('benchmark', 'complete'):
        raise ValueError(f"data.scene_set must be 'benchmark' or 'complete'; got {want!r}")
    with open(config['data']['split_file']) as f:
        have = json.load(f).get('scene_set', 'benchmark')
    if have != want:
        raise ValueError(f"data.scene_set is {want!r} but {config['data']['split_file']} declares scene_set {have!r}")
    return want


class MultiObjectDynamics(Dataset):
    """MOVi multi-object dataset with per-object vertices and per-frame padding+mask.

    Expects dataset directories like movi_* with frames frame_00000/... and files:
      - meshes/obj_i_faces.npy (per-object faces, static per dataset)
      - frame_xxxxx/obj_i_pos.npy (per-object vertices in world space)
    """

    def __init__(
        self,
        data_path: str,
        sequence_length: int = 3,
        stride: int = 1,
        split: str = 'train',
        train_ratio: float = 0.9,
        val_ratio: float = 0.1,
        split_file: Optional[str] = None,
        fast_index: bool = True,
        skip_first_frames: int = 0,
        skip_step_size: int = 1,
        max_vertices_per_object: Optional[int] = None,
        max_objects_per_frame: Optional[int] = None,
        anchor_k: Optional[int] = None,
        if_phys_param: bool = False,
        **kwargs
    ):
        self.data_path = Path(data_path)
        self.sequence_length = sequence_length
        self.stride = stride
        self.split = split
        self.skip_first_frames = skip_first_frames
        self.skip_step_size = skip_step_size
        self.max_vertices_per_object = max_vertices_per_object
        self.max_objects_per_frame = max_objects_per_frame
        self.if_phys_param = if_phys_param  # Strict mode for physics parameters

        # Point Cloud Mode configuration
        # When enabled, load point cloud files (obj_i_pc{N}_pos.npy) instead of mesh vertices
        if_point_cloud_config = kwargs.get('if_point_cloud', {})
        self.point_cloud_enable = if_point_cloud_config.get('enable', False)
        if self.point_cloud_enable:
            # the point-cloud path validates neither per-frame point counts nor per-resolution padding (a frame with a
            # missing point was accepted; resolutions above the first were truncated to it), and train.py does not
            # forward the setting: refused until implemented and validated
            raise NotImplementedError('point-cloud inputs (if_point_cloud.enable) are not supported by this loader')
        self.point_cloud_train_resolution = if_point_cloud_config.get('train_resolution', [512])
        self.point_cloud_eval_resolution = if_point_cloud_config.get('eval_resolution', [512])

        # Ensure resolutions are lists
        if isinstance(self.point_cloud_train_resolution, int):
            self.point_cloud_train_resolution = [self.point_cloud_train_resolution]
        if isinstance(self.point_cloud_eval_resolution, int):
            self.point_cloud_eval_resolution = [self.point_cloud_eval_resolution]

        # Select resolution based on split
        if self.point_cloud_enable:
            if split == 'train':
                self.point_cloud_resolutions = self.point_cloud_train_resolution
            else:
                self.point_cloud_resolutions = self.point_cloud_eval_resolution
            print(f"[{split}] Point Cloud Mode ENABLED:")
            print(f"  - Resolutions: {self.point_cloud_resolutions}")
            print(f"  - Data will be loaded from: frame_XXXXX/obj_i_pc{{N}}_pos.npy")
        else:
            self.point_cloud_resolutions = None

        # Anchor config (Nq): prefer explicit arg, else env ANCHOR_K, else default 32
        if anchor_k is None:
            env_anchor = os.environ.get('ANCHOR_K', '').strip()
            if env_anchor:
                try:
                    anchor_k = int(env_anchor)
                except Exception:
                    raise ValueError(f"Invalid ANCHOR_K environment value: '{env_anchor}' (expected int)")
            else:
                anchor_k = 32
        if anchor_k <= 0:
            raise ValueError(f"anchor_k must be positive; got {anchor_k}")
        self.anchor_k = int(anchor_k)

        # Object Permutation Augmentation (train only, at least 2 objects)
        augmentation_config = kwargs.get('augmentation', {})
        # The random-walk / independent noise blocks were never applied on this path (the dataset never stored them and
        # the helpers expect body1_/body2_ arrays, not MOVi's obj_positions): enabling them is an error, not a no-op.
        for _blk in ('random_walk', 'independent'):
            for _cfg in (augmentation_config, kwargs.get('noise_config') or {}):
                if ((_cfg or {}).get(_blk) or {}).get('enable', False):
                    raise NotImplementedError(f"data.augmentation.{_blk}.enable: true is not implemented for the MOVi loader "
                                              f"(it would be silently ignored); set it to false")
        # Anchor selection rule. 'index0' starts FPS at vertex 0 (order dependent);
        # 'centroid_farthest' starts at the vertex farthest from the centroid with lexicographic
        # tie-breaks, so the anchors depend on the vertex set only.
        self.anchor_fps_start = str(kwargs.get('anchor_fps_start', 'index0'))
        # Order-independent FPS ordering of every object's vertices (cached as obj_i_fpsorder_cf.npy);
        # the encoder's hierarchical levels take the first 50/25/12.5/... percent of this ordering.
        self.fps_order = bool(kwargs.get('fps_order', False))
        # data.object_scale_file: apply each object's metadata size to the stored vertices (src/data/object_scale.py)
        _scale_file = kwargs.get('object_scale_file')
        self._object_scale = ObjectScale(_scale_file) if _scale_file else None
        if self._object_scale is not None:
            # the folder name is compared without regard to case
            _ds_real = os.path.basename(os.path.realpath(str(data_path)))
            if (self._object_scale.converted or '').lower() != _ds_real.lower():
                raise ValueError(f"data.object_scale_file {_scale_file} was built for '{self._object_scale.converted}', "
                                 f"but the dataset path resolves to '{_ds_real}'")
        if self._object_scale is not None:
            print(f"[{split}] Object scale ACTIVE: vertices scaled about the object origin by the metadata size ({_scale_file})")
        # data.partial_view: partial point clouds of a simulated depth camera instead of the mesh vertices
        # (src/data/partial_view.py)
        _pv = kwargs.get('partial_view') or {}
        self._partial = None
        if _pv.get('enable', False):
            if self._object_scale is None:
                raise ValueError("data.partial_view needs data.object_scale_file for the objects' metadata sizes")
            self._partial = PartialView(_pv, Path(__file__).resolve().parents[2])
            if self._partial.min_points < self.anchor_k:
                raise ValueError(f'data.partial_view.min_points ({self._partial.min_points}) must be at least '
                                 f'anchor_k ({self.anchor_k})')
            self._partial_scenes = {}
            print(f"[{split}] {self._partial.describe(train=split == 'train')}")
        self._fps_order_by_dataset = {}
        if self.anchor_fps_start not in ('index0', 'centroid_farthest'):
            raise ValueError(f"data.anchor_fps_start must be 'index0' or 'centroid_farthest'; got '{self.anchor_fps_start}'")
        if self.anchor_fps_start != 'index0':
            print(f"[{split}] Anchor FPS start rule: {self.anchor_fps_start} (anchor cache suffix '_cf')")
        # Anchor resampling augmentation (train split only): with probability `prob` the FPS anchors
        # of every object are recomputed from a random start vertex.
        resample_config = augmentation_config.get('anchor_resampling', {})
        self.anchor_resample_prob = float(resample_config.get('prob', 0.5)) if resample_config.get('enable', False) else 0.0
        if self.anchor_resample_prob > 0 and split == 'train':
            print(f"[{split}] Anchor resampling augmentation enabled: prob={self.anchor_resample_prob:.2f}")
        permute_config = augmentation_config.get('object_permutation', {})
        self.permute_enabled = permute_config.get('enable', False) and split == 'train'
        self.permute_prob = permute_config.get('prob_apply', 0.5)
        self.permute_debug = permute_config.get('if_debug', False)
        if self.permute_enabled:
            print(f"[{split}] Object permutation augmentation enabled: prob={self.permute_prob:.2f}, debug={self.permute_debug}")

        # Load the dataset split file (required; splits are never auto-generated).
        if not split_file:
            raise ValueError(
                "split_file is required. Set 'data.split_file' in the config to a split JSON, "
                "e.g. one of the dataset_split_*.json files in the repository root."
            )

        if not os.path.exists(split_file):
            raise FileNotFoundError(
                f"Split file not found: {split_file}\n"
                "Check the 'data.split_file' path in the config. Split files "
                "(dataset_split_*.json) are provided in the repository root; new splits "
                "can be written for data produced by data_generation/movi_to_sim_data.py."
            )

        with open(split_file, 'r') as f:
            s = json.load(f)
        names = s['train_datasets'] if split == 'train' else s['val_datasets'] if split == 'val' else s['test_datasets']
        self.dataset_dirs = [self.data_path / n for n in names]
        if self._object_scale is not None:
            check_object_counts(self._object_scale, self.data_path, names)

        # Build sequences
        self.all_sequences = []  # each: { 'dataset': name, 'frames': [paths], 'frame_ids': [int] }

        # DDP-aware: only show progress bar on rank 0
        import torch.distributed as dist
        show_progress = True
        if dist.is_available() and dist.is_initialized():
            show_progress = (dist.get_rank() == 0)

        for ds in tqdm(self.dataset_dirs, desc=f"Indexing {split} datasets", unit="ds", disable=not show_progress):
            frames = scene_frame_dirs(ds)
            if self.skip_first_frames > 0:
                frames = frames[self.skip_first_frames:]
            if self.skip_step_size > 1:
                frames = frames[::self.skip_step_size]
            for start in range(0, len(frames) - sequence_length + 1, self.stride):
                seq = frames[start:start+sequence_length]
                if len(seq) == sequence_length:
                    self.all_sequences.append({
                        'dataset': ds.name,
                        'frames': [p.as_posix() for p in seq],
                        'frame_ids': [int(Path(p).name.split('_')[-1]) for p in seq]
                    })

        # Precompute/load per-object anchor indices and faces for each dataset from
        # the reference frame_00000; cached in memory.
        self._anchors_by_dataset = {}
        self._faces_by_dataset = {}

        # For point cloud mode, use the first resolution for __init__ validation
        init_pc_resolution = self.point_cloud_resolutions[0] if self.point_cloud_enable else None

        for ds in self.dataset_dirs:
            ds_name = ds.name
            frame0 = ds / 'frame_00000'
            if not (frame0.exists() or is_packed(ds)):
                raise FileNotFoundError(
                    f"Missing reference frame for dataset '{ds_name}': {frame0}. "
                    f"Expected a 'frame_00000' directory (or frames.npy) under {ds}.")

            # For point cloud mode, load point clouds for validation
            # For mesh mode, load mesh vertices
            ref_objs = self._load_frame_objects(frame0, pc_resolution=init_pc_resolution)
            if len(ref_objs) == 0:
                raise RuntimeError(f"Empty reference frame in {frame0}: no per-object vertex files found.")
            anchors_dir = ds / 'anchors'
            anchors_dir.mkdir(parents=True, exist_ok=True)
            ds_anchor_list = []
            ds_faces_list = []
            ds_order_list = []
            for i, V in enumerate(ref_objs):
                nv = V.shape[0]
                if self.max_vertices_per_object is not None and nv > self.max_vertices_per_object:
                    raise ValueError(
                        f"Dataset '{ds_name}' reference object {i} has {nv} vertices, "
                        f"exceeding max_vertices_per_object={self.max_vertices_per_object}.\n"
                        f"Either reduce vertex count in dataset or increase data.max_vertices_per_object in config."
                    )

                # Point cloud mode: anchors are unused, so use sequential placeholder
                # indices rather than mesh anchors (whose indices may be out of range).
                if self.point_cloud_enable:
                    idx = np.arange(min(self.anchor_k, nv), dtype=np.int32)
                    if len(idx) < self.anchor_k:
                        idx = np.pad(idx, (0, self.anchor_k - len(idx)), mode='constant')
                    ds_anchor_list.append(idx)
                    # Create empty faces (not used in point cloud mode)
                    ds_faces_list.append(np.zeros((0, 3), dtype=np.int64))
                    continue

                # Mesh vertex mode: standard anchor and faces loading.
                if nv < self.anchor_k:
                    raise ValueError(
                        f"Dataset '{ds_name}' object {i} has only {nv} vertices (< anchor_k={self.anchor_k}). "
                        f"Reduce ANCHOR_K or increase vertex cap (max_vertices_per_object).")
                _suffix = '' if self.anchor_fps_start == 'index0' else '_cf'
                anchor_path = anchors_dir / f"obj_{i}_anchors_fps{self.anchor_k}{_suffix}.npy"
                if anchor_path.exists():
                    idx = np.load(anchor_path)
                    if idx.ndim != 1 or idx.shape[0] != self.anchor_k:
                        raise ValueError(
                            f"Invalid anchor file shape: {anchor_path} (expected ({self.anchor_k},))")
                    if idx.dtype.kind not in ('i', 'u'):
                        raise ValueError(f"Invalid anchor file dtype: {anchor_path} (expected integer)")
                    if int(idx.max()) >= nv or int(idx.min()) < 0:
                        raise ValueError(
                            f"Anchor indices in {anchor_path} out of range [0,{nv-1}] for object {i}")
                    if np.unique(idx).shape[0] != idx.shape[0]:
                        raise ValueError(f"Anchor cache {anchor_path} repeats a vertex: {idx.tolist()}")
                    ds_anchor_list.append(idx.astype(np.int32, copy=False))
                else:
                    # Compute deterministic FPS anchors on the STORED (unscaled) frame-0 vertices, as
                    # pregen_cf_anchors.py does: FPS is invariant to the size scaling only in exact arithmetic (float32
                    # rounding of the scaled coordinates breaks the exact ties of symmetric meshes differently)
                    V_ref = load_frame(ds, 0)[i].astype(np.float32, copy=False)
                    idx = self._farthest_point_sample_np(V_ref, self.anchor_k, rule=self.anchor_fps_start)
                    # Atomic write: concurrent DDP ranks must never read a partially written cache.
                    tmp_path = anchor_path.with_name(f"{anchor_path.stem}.{os.getpid()}.tmp.npy")
                    np.save(tmp_path, idx.astype(np.int32, copy=False))
                    os.replace(tmp_path, anchor_path)
                    ds_anchor_list.append(idx.astype(np.int32, copy=False))
                if self.fps_order:
                    order_path = anchors_dir / f"obj_{i}_fpsorder_cf.npy"
                    if order_path.exists():
                        order = np.load(order_path)
                        if order.ndim != 1 or order.shape[0] != nv or order.dtype.kind not in ('i', 'u'):
                            raise ValueError(f"Invalid FPS-order cache {order_path}: shape {order.shape}, expected ({nv},)")
                        if not np.array_equal(np.sort(order), np.arange(nv)):
                            raise ValueError(f"FPS-order cache {order_path} is not a permutation of the {nv} vertices")
                    else:
                        # stored (unscaled) frame-0 vertices, as pregen_fps_order.py does (see the anchors above)
                        order = self._farthest_point_sample_np(load_frame(ds, 0)[i].astype(np.float32, copy=False),
                                                               nv, rule='centroid_farthest')
                        tmp_path = order_path.with_name(f"{order_path.stem}.{os.getpid()}.tmp.npy")
                        np.save(tmp_path, order.astype(np.int32, copy=False))
                        os.replace(tmp_path, order_path)
                    ds_order_list.append(order.astype(np.int32, copy=False))

                # Load faces for this object (static per dataset), and filter to nv cap
                faces_path = ds / 'meshes' / f'obj_{i}_faces.npy'
                if not faces_path.exists():
                    raise FileNotFoundError(
                        f"Missing faces file for dataset '{ds_name}' object {i}: {faces_path}")
                faces = np.load(faces_path)
                if faces.ndim != 2 or faces.shape[1] != 3:
                    raise ValueError(f"Invalid faces array at {faces_path}: shape={faces.shape}, expected (F,3)")
                # every face must index this object's vertices (the vertex count is never capped: an object above
                # max_vertices_per_object is an error), so an out-of-range face is corrupt data, not a face to drop
                if faces.dtype.kind not in ('i', 'u'):
                    raise ValueError(f"Faces of {faces_path} have dtype {faces.dtype}, expected integers")
                faces = faces.astype(np.int64, copy=False)
                if faces.size and (int(faces.min()) < 0 or int(faces.max()) >= nv):
                    raise ValueError(f"Faces of {faces_path} index vertices outside [0, {nv - 1}]")
                ds_faces_list.append(faces)
            self._anchors_by_dataset[ds_name] = ds_anchor_list
            self._faces_by_dataset[ds_name] = ds_faces_list
            if self.fps_order:
                self._fps_order_by_dataset[ds_name] = ds_order_list

        # Cache reference frames (frame_00000) per dataset to avoid repeated loading.
        self._reference_frames_cache = {}

        # Per-dataset caches: frame lists (avoids glob in __getitem__), frame IDs
        # (avoids string parsing per sample), and per-object/max vertex counts.
        self._all_frames_by_dataset = {}
        self._frame_id_cache = {}
        self._max_nv_by_dataset = {}
        self._mesh_nv_by_dataset = {}

        for ds in self.dataset_dirs:
            ds_name = ds.name
            all_frames = scene_frame_dirs(ds)
            self._all_frames_by_dataset[ds_name] = all_frames
            for fp in all_frames:
                frame_id = int(fp.name.split('_')[-1])  # "frame_00123" -> 123
                self._frame_id_cache[(ds_name, fp.name)] = frame_id
            # Frames are indexed by position: a missing or extra frame directory would silently change the time step
            # and misalign the ground truth (vis_traj.npy, which the evaluator indexes by frame number).
            _ids = [self._frame_id_cache[(ds_name, fp.name)] for fp in all_frames]
            _T_traj = np.load(ds / 'vis_traj.npy', mmap_mode='r').shape[1]
            if _ids != list(range(len(_ids))) or len(_ids) != _T_traj:
                raise ValueError(f"{ds}: frame directories must be frame_00000 .. frame_{_T_traj - 1:05d} without gaps "
                                 f"(vis_traj.npy has {_T_traj} frames); found {len(_ids)} directories, "
                                 f"first gap at {next((i for i, f in enumerate(_ids) if f != i), len(_ids))}")
            # Load frame_00000 once to record per-object vertex counts
            # (at point cloud resolution when enabled).
            frame0 = ds / 'frame_00000'
            ref_objs_data = self._load_frame_objects(frame0, pc_resolution=init_pc_resolution)
            nv_list = [obj.shape[0] for obj in ref_objs_data]
            self._mesh_nv_by_dataset[ds_name] = nv_list
            self._max_nv_by_dataset[ds_name] = max(nv_list) if nv_list else 0

        # Load physics parameters (mass, friction, restitution) per dataset.
        # With if_phys_param=True the file and every object entry must exist.
        self._physics_by_dataset = {}
        for ds in self.dataset_dirs:
            ds_name = ds.name
            physics_path = ds / 'physics_params.json'
            num_objs = len(self._anchors_by_dataset[ds_name])

            if self.if_phys_param:
                if not physics_path.exists():
                    raise FileNotFoundError(
                        f"Physics parameters required but file missing: {physics_path}\n"
                        f"Dataset: {ds_name} | if_phys_param=True requires physics_params.json\n"
                        f"Please generate physics parameters or set if_phys_param=False"
                    )

                # Parse errors propagate deliberately.
                with open(physics_path, 'r') as f:
                    physics_params = json.load(f)

                phys_list = []
                for i in range(num_objs):
                    obj_key = f"obj_{i}"
                    if obj_key not in physics_params:
                        raise KeyError(
                            f"Missing physics parameters for {obj_key} in {physics_path}\n"
                            f"Dataset: {ds_name} | Found objects: {list(physics_params.keys())}\n"
                            f"Required: obj_0 through obj_{num_objs-1}"
                        )

                    p = physics_params[obj_key]
                    if 'mass' not in p:
                        raise KeyError(f"Missing 'mass' field for {obj_key} in {physics_path}")
                    if 'friction' not in p:
                        raise KeyError(f"Missing 'friction' field for {obj_key} in {physics_path}")
                    if 'restitution' not in p:
                        raise KeyError(f"Missing 'restitution' field for {obj_key} in {physics_path}")

                    phys_list.append([
                        float(p['mass']),
                        float(p['friction']),
                        float(p['restitution'])
                    ])

                self._physics_by_dataset[ds_name] = np.array(phys_list, dtype=np.float32)  # (Nobj, 3)
            else:
                # Permissive mode: substitute defaults for missing entries.
                if physics_path.exists():
                    with open(physics_path, 'r') as f:
                        physics_params = json.load(f)
                    phys_list = []
                    for i in range(num_objs):
                        obj_key = f"obj_{i}"
                        if obj_key in physics_params:
                            p = physics_params[obj_key]
                            phys_list.append([
                                float(p.get('mass', 1.0)),
                                float(p.get('friction', 0.5)),
                                float(p.get('restitution', 0.5))
                            ])
                        else:
                            phys_list.append([1.0, 0.5, 0.5])
                    self._physics_by_dataset[ds_name] = np.array(phys_list, dtype=np.float32)
                else:
                    # File missing entirely: default values for all objects.
                    self._physics_by_dataset[ds_name] = np.array([[1.0, 0.5, 0.5]] * num_objs, dtype=np.float32)

        # Multi-timescale training configuration; step sizes are expressed in
        # normalized time units.
        random_step_size_config = kwargs.get('random_step_size', {})
        self.random_step_size_enable = random_step_size_config.get('enable', False)
        self.random_step_sizes = random_step_size_config.get('step_sizes', [1, 2, 5, 10])
        self.random_start_frame = random_step_size_config.get('random_start_frame', True)
        self.sampling_weights = random_step_size_config.get('sampling_weights', None)

        if self.random_step_size_enable:
            if not isinstance(self.random_step_sizes, list) or len(self.random_step_sizes) == 0:
                raise ValueError(
                    f"random_step_size.step_sizes must be a non-empty list, got: {self.random_step_sizes}"
                )

            for ss in self.random_step_sizes:
                if not isinstance(ss, int) or ss <= 0:
                    raise ValueError(
                        f"All step_sizes must be positive integers, got: {self.random_step_sizes}"
                    )

            if self.sampling_weights is None:
                # Default to uniform sampling.
                self.sampling_weights = [1.0 / len(self.random_step_sizes)] * len(self.random_step_sizes)
            else:
                if len(self.sampling_weights) != len(self.random_step_sizes):
                    raise ValueError(
                        f"sampling_weights length ({len(self.sampling_weights)}) must match "
                        f"step_sizes length ({len(self.random_step_sizes)})"
                    )

                weights_sum = sum(self.sampling_weights)
                if abs(weights_sum - 1.0) > 1e-6:
                    raise ValueError(
                        f"sampling_weights must sum to 1.0, got sum={weights_sum:.6f}\n"
                        f"Weights: {self.sampling_weights}"
                    )

            print(f"\n{'='*60}")
            print(f"[{split}] Multi-timescale training ENABLED:")
            print(f"  - step_sizes: {self.random_step_sizes}")
            print(f"  - sampling_weights: {self.sampling_weights}")
            print(f"  - random_start_frame: {self.random_start_frame}")
            print(f"  - if_debug: {random_step_size_config.get('if_debug', False)}")
            print(f"{'='*60}\n")

            # Every dataset must have enough frames for every configured step size.
            print(f"[{split}] Validating dataset frame counts for multi-timescale training...")
            insufficient_datasets = []
            for ds in self.dataset_dirs:
                frames = scene_frame_dirs(ds)
                if self.skip_first_frames > 0:
                    frames = frames[self.skip_first_frames:]
                available_frames = len(frames)

                for step_size in self.random_step_sizes:
                    frames_needed = (self.sequence_length - 1) * step_size + 1
                    if available_frames < frames_needed:
                        insufficient_datasets.append({
                            'dataset': ds.name,
                            'available': available_frames,
                            'needed': frames_needed,
                            'step_size': step_size,
                            'sequence_length': self.sequence_length
                        })

            if insufficient_datasets:
                error_msg = f"[{split}] Frame validation failed: {len(insufficient_datasets)} datasets have insufficient frames:\n"
                for info in insufficient_datasets[:5]:
                    error_msg += (
                        f"  - {info['dataset']}: available={info['available']}, "
                        f"needed={info['needed']} (step_size={info['step_size']}, "
                        f"sequence_length={info['sequence_length']})\n"
                    )
                if len(insufficient_datasets) > 5:
                    error_msg += f"  ... and {len(insufficient_datasets) - 5} more datasets\n"
                error_msg += (
                    f"\nSolutions:\n"
                    f"  1. Reduce max step_size in config: data.random_step_size.step_sizes\n"
                    f"  2. Reduce sequence_length in config: data.sequence_length\n"
                    f"  3. Filter out datasets with insufficient frames\n"
                    f"  4. Reduce skip_first_frames in config: data.skip_first_frames (currently {self.skip_first_frames})"
                )
                raise ValueError(error_msg)

            print(f"[{split}] All {len(self.dataset_dirs)} datasets have sufficient frames for step_sizes {self.random_step_sizes}")
        else:
            # Disabled: use the fixed skip_step_size
            print(f"[{split}] Multi-timescale training disabled, using skip_step_size={self.skip_step_size}")

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

    def get_cache_stats(self) -> dict:
        """Return cache statistics for monitoring."""
        return {
            'numpy_cache_info': _cached_numpy_load.cache_info()._asdict(),
            'reference_frames_cached': len(self._reference_frames_cache),
            'reference_frames_datasets': list(self._reference_frames_cache.keys())
        }

    @staticmethod
    def clear_numpy_cache():
        """Clear the global numpy LRU cache to free memory."""
        _cached_numpy_load.cache_clear()
        logger.info("Cleared numpy LRU cache")

    def _load_frame_objects(self, frame_dir: Path, pc_resolution: Optional[int] = None) -> List['np.ndarray']:
        """
        Load per-object vertex/point data for a single frame.

        Args:
            frame_dir: Path to frame directory
            pc_resolution: If provided, load point cloud at this resolution instead of mesh vertices

        Returns:
            List of (Ni, 3) arrays, one per object
        """
        # Data exceeding the configured limits is rejected rather than truncated.
        objs = []
        idx = 0

        if pc_resolution is not None:
            if self._object_scale is not None:
                raise NotImplementedError('data.object_scale_file is not supported with point-cloud inputs')
            # Point cloud mode: load from obj_i_pc{N}_pos.npy
            while (frame_dir / f'obj_{idx}_pc{pc_resolution}_pos.npy').exists():
                pc_file = frame_dir / f'obj_{idx}_pc{pc_resolution}_pos.npy'
                V = cached_load_npy(pc_file)  # (N_points, 3)

                # Point clouds have a fixed size; no vertex limit check needed.
                objs.append(V.astype('float32', copy=False))
                idx += 1

            # Mesh vertices exist but the requested point cloud resolution does not.
            if idx == 0 and (frame_dir / 'obj_0_pos.npy').exists():
                raise FileNotFoundError(
                    f"Point cloud files not found at resolution {pc_resolution} in {frame_dir}.\n"
                    f"Expected: obj_i_pc{pc_resolution}_pos.npy\n"
                    f"Mesh vertices exist but point cloud mode is enabled.\n"
                    f"Please use a dataset with pre-computed point clouds."
                )
        else:
            # Mesh vertex mode: load from obj_i_pos.npy (packed layout: the same arrays from frames.npy)
            packed = (load_frame(frame_dir.parent, int(frame_dir.name.split('_')[-1]))
                      if is_packed(frame_dir.parent) else None)
            while (idx < len(packed)) if packed is not None else (frame_dir / f'obj_{idx}_pos.npy').exists():
                V = packed[idx] if packed is not None else cached_load_npy(frame_dir / f'obj_{idx}_pos.npy')  # (Ni, 3)
                if self._object_scale is not None:
                    V = self._object_scale.apply(frame_dir / f'obj_{idx}_pos.npy', V)

                if self.max_vertices_per_object is not None and V.shape[0] > self.max_vertices_per_object:
                    raise ValueError(
                        f"Object {idx} in {frame_dir} has {V.shape[0]} vertices, "
                        f"exceeding max_vertices_per_object={self.max_vertices_per_object}.\n"
                        f"Either reduce vertex count in dataset or increase data.max_vertices_per_object in config."
                    )

                objs.append(V.astype('float32', copy=False))
                idx += 1

        if self.max_objects_per_frame is not None and len(objs) > self.max_objects_per_frame:
            raise ValueError(
                f"Frame {frame_dir} has {len(objs)} objects, "
                f"exceeding max_objects_per_frame={self.max_objects_per_frame}.\n"
                f"Either reduce object count in dataset or increase data.max_objects_per_frame in config."
            )

        return objs

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        info = self.all_sequences[idx]

        # Point Cloud Mode: Sample resolution for this sample
        pc_resolution = None
        if self.point_cloud_enable:
            pc_resolution = np.random.choice(self.point_cloud_resolutions)

        # Multi-timescale training: sample a step size for this sequence.
        if self.random_step_size_enable:
            step_size = np.random.choice(self.random_step_sizes, p=self.sampling_weights)

            ds_name = info['dataset']
            all_frames = self._all_frames_by_dataset[ds_name]

            # Apply skip_first_frames
            if self.skip_first_frames > 0:
                all_frames = all_frames[self.skip_first_frames:]

            # Total frames spanned by the sequence, e.g. sequence_length=4 with
            # step_size=10 uses frames [0, 10, 20, 30] and spans 31 frames.
            student_reach = (self.sequence_length - 1) * step_size
            total_frames_needed = student_reach + 1

            # Frame counts were validated in __init__; assert as a sanity check.
            assert len(all_frames) >= total_frames_needed, \
                f"Unexpected: insufficient frames in {ds_name} (should have been caught in __init__)"

            # Select a start frame within the valid range.
            if self.random_start_frame:
                max_start_frame = len(all_frames) - total_frames_needed
                start_frame = np.random.randint(0, max_start_frame + 1)
            else:
                # Deterministic start frame derived from idx (for reproducibility).
                max_start_frame = len(all_frames) - total_frames_needed
                start_frame = idx % (max_start_frame + 1)

            # Frame sequence [start_frame, start_frame + step_size, ...].
            frame_indices = [start_frame + i * step_size for i in range(self.sequence_length)]
            frame_paths = [all_frames[i] for i in frame_indices]
            frame_ids = np.array([self._frame_id_cache[(ds_name, fp.name)] for fp in frame_paths], dtype=np.int32)

            selected_step_size = step_size
        else:
            # Fixed step size: use pre-built sequences with skip_step_size.
            ds_name = info['dataset']
            frame_paths = [Path(p) for p in info['frames']]
            frame_ids = np.asarray(info.get('frame_ids', []), dtype=np.int32)
            selected_step_size = self.skip_step_size

        if self._partial is not None:
            return self._partial_item(ds_name, frame_ids, selected_step_size)

        # Load per-frame objects (ragged); point clouds are loaded instead of mesh
        # vertices when pc_resolution is set.
        frames_objs: List[List[np.ndarray]] = [self._load_frame_objects(fp, pc_resolution=pc_resolution) for fp in frame_paths]
        if pc_resolution is None:
            # every frame holds the reference frame's objects with the same vertex counts (a missing or truncated file
            # would otherwise become padding or a shorter object without notice)
            _ref_nv = self._mesh_nv_by_dataset[ds_name]
            for fp, fo in zip(frame_paths, frames_objs):
                if [V.shape[0] for V in fo] != _ref_nv:
                    raise ValueError(f"{fp}: per-object vertex counts {[V.shape[0] for V in fo]} differ from frame_00000 "
                                     f"{_ref_nv}")
        num_objs = max(len(f) for f in frames_objs)

        if self.max_vertices_per_object is not None:
            max_nv = int(self.max_vertices_per_object)
        else:
            max_nv = self._max_nv_by_dataset[ds_name]

        T = len(frames_objs)

        # Allocate padded arrays.
        obj_positions = np.full((T, num_objs, max_nv, 3), PADDING_VALUE, dtype='float32')
        valid_obj_mask = np.zeros((num_objs,), dtype='bool')
        valid_vtx_mask = np.zeros((T, num_objs, max_nv), dtype='bool')

        for t, fo in enumerate(frames_objs):
            for i, V in enumerate(fo):
                nv = min(V.shape[0], max_nv)
                obj_positions[t, i, :nv, :] = V[:nv]
                valid_obj_mask[i] = True
                valid_vtx_mask[t, i, :nv] = True

        # Reference pose is always the first frame of the dataset (global frame_00000),
        # cached per dataset. Point cloud mode includes the resolution in the cache key.
        cache_key = f"{ds_name}_pc{pc_resolution}" if pc_resolution else ds_name
        if cache_key not in self._reference_frames_cache:
            ds_root = self.data_path / ds_name
            frame0 = ds_root / 'frame_00000'
            if not (frame0.exists() or is_packed(ds_root)):
                raise FileNotFoundError(
                    f"Missing reference frame for dataset '{ds_name}': {frame0}. "
                    f"Expected a 'frame_00000' directory (or frames.npy) under {ds_root}."
                )
            ref_objs = self._load_frame_objects(frame0, pc_resolution=pc_resolution)
            if len(ref_objs) == 0:
                raise RuntimeError(
                    f"Empty reference frame in {frame0}: no per-object vertex files found."
                )
            self._reference_frames_cache[cache_key] = ref_objs
        else:
            ref_objs = self._reference_frames_cache[cache_key]
        # Every frame in the sequence must contain all reference objects, each stored
        # as a non-empty (N, 2|3) array. Point cloud mode is already validated inside
        # _load_frame_objects, so it is skipped here, and so is the packed layout (frames.npy holds every object of
        # every frame by construction; scene_frames checks its shape against frames_counts.npy).
        expected_num_objs = len(ref_objs)
        if pc_resolution is None and not is_packed(self.data_path / ds_name):
            for t, fdir in enumerate(frame_paths):
                for i_ref in range(expected_num_objs):
                    npy_path = fdir / f"obj_{i_ref}_pos.npy"
                    if not npy_path.exists():
                        raise RuntimeError(
                            f"Missing expected file for dataset '{ds_name}': {npy_path} "
                            f"(t={t}, obj={i_ref}). Ensure all frames contain the same set of objects as frame_00000.")
                    try:
                        arr_chk = np.load(npy_path)
                    except Exception as e:
                        raise RuntimeError(
                            f"Failed to load file: {npy_path} (t={t}, obj={i_ref}). Error: {e}")
                    if not (arr_chk.ndim == 2 and arr_chk.shape[1] in (2, 3) and arr_chk.shape[0] > 0):
                        raise RuntimeError(
                            f"Invalid/empty vertices at {npy_path}: shape={arr_chk.shape} "
                            f"(t={t}, obj={i_ref}).")
        # Build the reference array with fixed Nv (pad/trim to max_nv).
        ref_obj_positions = np.full((num_objs, max_nv, 3), PADDING_VALUE, dtype='float32')
        for i, V in enumerate(ref_objs[:num_objs]):
            nv = min(V.shape[0], max_nv)
            ref_obj_positions[i, :nv, :] = V[:nv]

        # Load anchors (per object) for this dataset (computed in __init__)
        if ds_name not in self._anchors_by_dataset:
            raise KeyError(f"Missing anchors for dataset '{ds_name}'. Initialization may have failed.")
        ds_anchor_list = self._anchors_by_dataset[ds_name]
        if len(ds_anchor_list) < num_objs:
            raise ValueError(
                f"Anchors count ({len(ds_anchor_list)}) < objects in sequence ({num_objs}) for dataset '{ds_name}'.")
        anchor_idx = np.zeros((num_objs, self.anchor_k), dtype=np.int32)
        for i in range(num_objs):
            idx_i = ds_anchor_list[i]
            if idx_i.shape[0] != self.anchor_k:
                raise ValueError(
                    f"Anchor length mismatch for dataset '{ds_name}' obj {i}: got {idx_i.shape[0]}, expected {self.anchor_k}")
            anchor_idx[i] = idx_i
        vtx_rank = None
        if self.fps_order:
            # rank[i, v] = position of vertex v in object i's FPS ordering (padding: max_nv)
            vtx_rank = np.full((num_objs, max_nv), max_nv, dtype=np.int32)
            ds_order_list = self._fps_order_by_dataset[ds_name]
            for i in range(num_objs):
                order_i = ds_order_list[i]
                nv_i = min(order_i.shape[0], max_nv)
                vtx_rank[i, order_i[:nv_i]] = np.arange(nv_i, dtype=np.int32)
        if (self.split == 'train' and self.anchor_resample_prob > 0 and not self.point_cloud_enable
                and np.random.rand() < self.anchor_resample_prob):
            _before = anchor_idx.copy()
            for i in range(num_objs):
                nv_i = min(ref_objs[i].shape[0], max_nv)
                if nv_i >= self.anchor_k:
                    anchor_idx[i] = self._farthest_point_sample_np(
                        ref_objs[i][:nv_i].astype(np.float32, copy=False), self.anchor_k,
                        rule=self.anchor_fps_start, start=int(np.random.randint(nv_i)))
            if not getattr(self, '_anchor_resample_logged', False):
                self._anchor_resample_logged = True
                print(f"[train] Anchor resampling APPLIED (pid {os.getpid()}): obj0 {_before[0].tolist()} -> {anchor_idx[0].tolist()} "
                      f"(rule={self.anchor_fps_start})", flush=True)

        # Object permutation must be applied before faces are gathered so that
        # faces_list[i] matches the permuted obj_i. Only faces for the current
        # sample's objects are taken.
        faces_list = self._faces_by_dataset[ds_name][:num_objs]
        if self.permute_debug:
            print(f"[object permutation] dataset={ds_name}, num_objs={num_objs}, faces_list length={len(faces_list)}, total_faces={len(self._faces_by_dataset[ds_name])}")
        perm = None  # Set only if permutation is applied
        if self.permute_enabled and np.random.random() < self.permute_prob:
            if num_objs >= 2:
                # Shuffle all objects; there is no ground mesh to exclude.
                perm = np.arange(num_objs)
                np.random.shuffle(perm)

                if self.permute_debug:
                    print(f"[object permutation] dataset={ds_name}, idx={idx}, num_objs={num_objs}")
                    print(f"  Original order: {list(range(num_objs))}")
                    print(f"  Shuffled order: {perm.tolist()}")
                    print(f"  Mapping: {' '.join([f'{i}->{perm[i]}' for i in range(num_objs)])}")

                # Permute all object-indexed data.
                obj_positions = obj_positions[:, perm]  # (T, Nobj, Nv, 3)
                ref_obj_positions = ref_obj_positions[perm]  # (Nobj, Nv, 3)
                ref_objs = [ref_objs[i] for i in perm]  # Permute ref_objs list
                # Permute frames_objs (list of lists)
                frames_objs = [[fo[i] for i in perm] for fo in frames_objs]
                valid_obj_mask = valid_obj_mask[perm]  # (Nobj,)
                valid_vtx_mask = valid_vtx_mask[:, perm]  # (T, Nobj, Nv)
                anchor_idx = anchor_idx[perm]  # (Nobj, Nq)
                if vtx_rank is not None:
                    vtx_rank = vtx_rank[perm]  # (Nobj, Nv)
                # Permute faces_list
                faces_list = [faces_list[i] for i in perm]

        # time_step_id: first global frame id of this sequence (scalar)
        ts_id_val = int(frame_ids[0]) if frame_ids.size > 0 else -1

        # Get physics parameters for this dataset (permuted if augmentation applied)
        phys_params = self._physics_by_dataset[ds_name]  # (Nobj_full, 3)
        phys_params = phys_params[:num_objs]  # Trim to current sample's object count
        if perm is not None:
            phys_params = phys_params[perm]  # Apply permutation to match permuted objects

        sample = {
            'obj_positions': torch.from_numpy(obj_positions),
            'valid_obj_mask': torch.from_numpy(valid_obj_mask),
            'valid_vtx_mask': torch.from_numpy(valid_vtx_mask),
            'ref_obj_positions': torch.from_numpy(ref_obj_positions),
            'dt': torch.tensor(1.0, dtype=torch.float32),
            'frame_ids': torch.from_numpy(frame_ids) if frame_ids.size > 0 else torch.tensor([], dtype=torch.int32),
            'time_step_id': torch.tensor(ts_id_val, dtype=torch.int64),
            'anchor_idx': torch.from_numpy(anchor_idx),
            **({'vtx_rank': torch.from_numpy(vtx_rank)} if vtx_rank is not None else {}),
            'dataset_name': ds_name,
            'faces_list': faces_list,  # Store permuted faces (list of numpy arrays)
            'permutation': perm if perm is not None else np.arange(num_objs),  # Store permutation order (identity if not permuted)
            'phys_params': torch.from_numpy(phys_params),  # (Nobj, 3): [mass, friction, restitution]
            'skip_step_size': torch.tensor(selected_step_size, dtype=torch.float32),  # Time step size for temporal encoding (randomly sampled when multi-timescale enabled)
            'is_point_cloud': pc_resolution is not None,  # Whether this sample uses point cloud data
        }

        if pc_resolution is not None:
            sample['point_cloud_resolution'] = pc_resolution


        # Apply noise augmentation (train split only)
        if getattr(self, 'split', 'train') == 'train' and getattr(self, 'noise_config', None):
            rw_cfg = self.noise_config.get('random_walk', {})
            if rw_cfg.get('enable', False):
                start_epoch = rw_cfg.get('start_epoch', 0)
                cur_epoch = getattr(self, 'current_epoch', 0)
                if cur_epoch >= start_epoch:
                    sample, noise_info = apply_random_walk_noise_to_batch(
                        sample, self.noise_config, is_training=True
                    )
                    sample['noise_info'] = noise_info
            elif self.noise_config.get('independent', {}).get('enable', False):
                sample = apply_independent_noise(sample, self.noise_config, is_training=True)

        return sample

    def _partial_scene(self, ds_name: str) -> dict:
        """Shapes, metadata sizes and pose trajectory (T, N, 7: position, quaternion w x y z) of a scene."""
        info = self._partial_scenes.get(ds_name)
        if info is None:
            ds = self.data_path / ds_name
            shapes = [o['asset_id'] for o in json.loads((ds / 'conversion.json').read_text())['objects']]
            traj = np.load(ds / 'vis_traj.npy')[0]
            sizes = self._object_scale.scales[ds_name]
            if not (len(shapes) == len(sizes) == traj.shape[1]):
                raise ValueError(f'{ds}: {len(shapes)} shapes, {len(sizes)} sizes, {traj.shape[1]} trajectories')
            info = {'shapes': shapes, 'sizes': np.asarray(sizes, dtype=np.float64), 'traj': traj}
            self._partial_scenes[ds_name] = info
        return info

    def _partial_item(self, ds_name: str, frame_ids: np.ndarray, step_size: int) -> Dict[str, torch.Tensor]:
        """Sample whose objects are the partial point clouds a depth camera observes at sample frame
        partial_view.observe_index, moved with the ground-truth poses to every frame (src/data/partial_view.py).

        Objects with fewer than min_points observed points are left out; `partial_obj_index` maps the sample's objects
        to the scene's. Points are in farthest-point order, so the anchors are the first anchor_k points and the FPS
        rank of a point is its index. Evaluation samples also carry the index of the observation frame in the sample
        (`partial_obs_index`) and the full collision meshes of all objects at that frame (`partial_full_obs`) for
        scoring.
        """
        info = self._partial_scene(ds_name)
        traj = info['traj']
        fids = np.asarray(frame_ids, dtype=np.int64)
        T = len(fids)
        k = min(self._partial.observe_index, T - 1)
        t_obs = int(fids[k])
        train = self.split == 'train'
        rng = np.random.default_rng(None if train else zlib.crc32(f'{ds_name}:{t_obs}'.encode()))
        pos_obs = traj[t_obs, :, :3].astype(np.float64)
        quat_obs = traj[t_obs, :, 3:7].astype(np.float64)
        # The model's contact search needs at least two objects. A training view that observes fewer is redrawn; after
        # 8 such views the evaluation camera, which observes nearly all objects, is used instead, or the whole objects
        # for a scene without a recorded camera (augmented training scenes).
        for attempt in range(9 if train else 1):
            use_train_camera = train and attempt < 8
            if train and not use_train_camera and ds_name not in self._partial.cameras:
                local = self._partial.whole_objects(info['shapes'])
            else:
                local = self._partial.observe(ds_name, info['shapes'], info['sizes'], pos_obs, quat_obs,
                                              use_train_camera, rng)
            observed = np.array([i for i, u in enumerate(local) if u is not None], dtype=np.int64)
            if observed.size >= 2:
                break
        if observed.size < 2:
            raise RuntimeError(f'{ds_name}: {observed.size} object(s) observed at frame {t_obs}, at least 2 needed '
                               f'({self._partial.describe(train=False)})')

        n_obj = int(observed.size)
        nmax = max(local[i].shape[0] for i in observed)
        P = traj[fids][:, observed, :3].astype(np.float64)                 # (T, n_obj, 3)
        Rq = quat_to_rotmat(traj[fids][:, observed, 3:7])                  # (T, n_obj, 3, 3)
        P0 = traj[0, observed, :3].astype(np.float64)                      # reference pose: global frame 0
        R0 = quat_to_rotmat(traj[0, observed, 3:7])
        obj_positions = np.full((T, n_obj, nmax, 3), PADDING_VALUE, dtype='float32')
        valid_vtx_mask = np.zeros((T, n_obj, nmax), dtype='bool')
        ref_obj_positions = np.full((n_obj, nmax, 3), PADDING_VALUE, dtype='float32')
        vtx_rank = np.full((n_obj, nmax), nmax, dtype=np.int32)
        for j, i in enumerate(observed):
            U = local[i] * info['sizes'][i]                                 # object frame, metres
            n = U.shape[0]
            obj_positions[:, j, :n] = P[:, j, None, :] + np.einsum('tab,nb->tna', Rq[:, j], U)
            valid_vtx_mask[:, j, :n] = True
            ref_obj_positions[j, :n] = P0[j] + U @ R0[j].T
            vtx_rank[j, :n] = np.arange(n, dtype=np.int32)
        anchor_idx = np.tile(np.arange(self.anchor_k, dtype=np.int32), (n_obj, 1))
        valid_obj_mask = np.ones((n_obj,), dtype='bool')
        phys_params = self._physics_by_dataset[ds_name][observed]
        faces_list = [np.zeros((0, 3), dtype=np.int64) for _ in range(n_obj)]

        perm = None
        if self.permute_enabled and np.random.random() < self.permute_prob and n_obj >= 2:
            perm = np.random.permutation(n_obj)
            obj_positions = obj_positions[:, perm]
            valid_vtx_mask = valid_vtx_mask[:, perm]
            ref_obj_positions = ref_obj_positions[perm]
            vtx_rank = vtx_rank[perm]
            anchor_idx = anchor_idx[perm]
            phys_params = phys_params[perm]
            observed = observed[perm]
            faces_list = [faces_list[p] for p in perm]

        sample = {
            'obj_positions': torch.from_numpy(obj_positions),
            'valid_obj_mask': torch.from_numpy(valid_obj_mask),
            'valid_vtx_mask': torch.from_numpy(valid_vtx_mask),
            'ref_obj_positions': torch.from_numpy(ref_obj_positions),
            'dt': torch.tensor(1.0, dtype=torch.float32),
            'frame_ids': torch.from_numpy(fids.astype(np.int32)),
            'time_step_id': torch.tensor(int(fids[0]), dtype=torch.int64),
            'anchor_idx': torch.from_numpy(anchor_idx),
            **({'vtx_rank': torch.from_numpy(vtx_rank)} if self.fps_order else {}),
            'dataset_name': ds_name,
            'faces_list': faces_list,
            'permutation': perm if perm is not None else np.arange(n_obj),
            'phys_params': torch.from_numpy(np.ascontiguousarray(phys_params)),
            'skip_step_size': torch.tensor(step_size, dtype=torch.float32),
            'is_point_cloud': False,
            'partial_obj_index': observed,
        }
        if not train:
            sample['partial_obs_index'] = k
            sample['partial_full_obs'] = self._load_frame_objects(self.data_path / ds_name / f'frame_{t_obs:05d}')
        return sample

    def set_epoch(self, epoch: int):
        """Set current epoch for noise scheduling."""
        self.current_epoch = int(epoch)

    @staticmethod
    def _farthest_point_sample_np(vertices: 'np.ndarray', k: int, rule: str = 'index0',
                                  start: Optional[int] = None) -> 'np.ndarray':
        """Deterministic farthest point sampling on an (N,3) array. Returns indices of shape (k,).

        rule='index0': start at vertex 0, ties resolved by the lowest index.
        rule='centroid_farthest': start at the vertex farthest from the centroid and resolve every
        tie by the lexicographically smallest coordinates, so the result depends on the vertex set
        only, not on its ordering. An explicit `start` overrides the rule's start vertex.
        """
        if vertices.ndim != 2 or vertices.shape[1] != 3:
            raise ValueError(f"Expected vertices shape (N,3); got {vertices.shape}")
        n = vertices.shape[0]
        if k > n:
            raise ValueError(f"FPS requested k={k} but only N={n} vertices available")
        V = vertices.astype(np.float64, copy=False)
        lex = rule == 'centroid_farthest'

        def argmax(d):
            if not lex:
                return int(np.argmax(d))
            cand = np.flatnonzero(d >= d.max() - 1e-9 * max(1.0, abs(float(d.max()))))
            if cand.shape[0] == 1:
                return int(cand[0])
            order = np.lexsort((V[cand, 2], V[cand, 1], V[cand, 0]))
            return int(cand[order[0]])

        selected = np.empty((k,), dtype=np.int64)
        dists = np.full((n,), np.inf, dtype=np.float64)
        if start is not None:
            selected[0] = int(start)
        elif lex:
            selected[0] = argmax(((V - V.mean(axis=0)) ** 2).sum(axis=1))
        else:
            selected[0] = 0
        last = V[selected[0]]
        for i in range(1, k):
            delta = V - last
            dists = np.minimum(dists, (delta * delta).sum(axis=1))
            # Never re-select a vertex: meshes may contain coincident duplicates, whose distance to an
            # already selected vertex is exactly zero. Excluding selected indices keeps the output a
            # permutation; duplicates are appended last (ties resolved by the rule above).
            dists[selected[:i]] = -np.inf
            idx = argmax(dists)
            selected[i] = idx
            last = V[idx]
        return selected.astype(np.int32, copy=False)


def make_collate_fn(max_objects_per_frame: Optional[int] = None, max_vertices_per_object: Optional[int] = None):
    """Factory function to create a collate_fn with fixed padding dimensions.

    Args:
        max_objects_per_frame: Fixed maximum number of objects per frame (None = dynamic)
        max_vertices_per_object: Fixed maximum vertices per object (None = dynamic)

    Returns:
        Configured collate function
    """
    def collate_fn(batch: List[Dict]) -> Dict[str, torch.Tensor]:
        """Custom collate function for batching sequences."""
        # Detect multi-object sample: obj_positions present with per-sample padding
        if 'obj_positions' in batch[0]:
            # Determine dimensions
            T = batch[0]['obj_positions'].shape[0]

            # Use fixed padding if configured, otherwise compute from batch
            if max_objects_per_frame is not None:
                max_nobj = int(max_objects_per_frame)
                # Validate that no sample exceeds the fixed limit
                for s in batch:
                    _, nobj, _, _ = s['obj_positions'].shape
                    if nobj > max_nobj:
                        raise ValueError(
                            f"Sample has {nobj} objects, exceeding max_objects_per_frame={max_nobj}. "
                            f"Increase data.max_objects_per_frame in config or reduce object count in dataset."
                        )
            else:
                # Dynamic padding
                max_nobj = 0
                for s in batch:
                    _, nobj, _, _ = s['obj_positions'].shape
                    if nobj > max_nobj:
                        max_nobj = nobj

            if max_vertices_per_object is not None:
                max_nv = int(max_vertices_per_object)
                # Validate that no sample exceeds the fixed limit
                for s in batch:
                    _, _, nv, _ = s['obj_positions'].shape
                    if nv > max_nv:
                        raise ValueError(
                            f"Sample has {nv} vertices, exceeding max_vertices_per_object={max_nv}. "
                            f"Increase data.max_vertices_per_object in config or reduce vertex count in dataset."
                        )
            else:
                # Dynamic padding
                max_nv = 0
                for s in batch:
                    _, _, nv, _ = s['obj_positions'].shape
                    if nv > max_nv:
                        max_nv = nv

            B = len(batch)
            device = batch[0]['obj_positions'].device

            # Allocate padded tensors.
            obj_positions = torch.full((B, T, max_nobj, max_nv, 3), PADDING_VALUE, dtype=torch.float32, device=device)
            valid_obj_mask = torch.zeros((B, max_nobj), dtype=torch.bool, device=device)
            valid_vtx_mask = torch.zeros((B, T, max_nobj, max_nv), dtype=torch.bool, device=device)
            ref_obj_positions = torch.full((B, max_nobj, max_nv, 3), PADDING_VALUE, dtype=torch.float32, device=device)
            # Optional anchors, padded per object to max_nobj with fixed Nq.
            # anchor_idx pads with index 0; anchor_mask pads with False.
            has_anchors = 'anchor_idx' in batch[0]
            has_rank = 'vtx_rank' in batch[0]
            if has_rank:
                vtx_rank = torch.full((B, max_nobj, max_nv), int(max_nv), dtype=torch.int64, device=device)
            if has_anchors:
                anchor_k = int(batch[0]['anchor_idx'].shape[1])
                anchor_idx = torch.zeros((B, max_nobj, anchor_k), dtype=torch.int64, device=device)
                anchor_mask = torch.zeros((B, max_nobj, anchor_k), dtype=torch.bool, device=device)

            # Fill per-sample
            for b, s in enumerate(batch):
                t_, nobj, nv, _ = s['obj_positions'].shape
                obj_positions[b, :, :nobj, :nv, :] = s['obj_positions']
                valid_obj_mask[b, :nobj] = s['valid_obj_mask']
                valid_vtx_mask[b, :, :nobj, :nv] = s['valid_vtx_mask']
                ref_obj_positions[b, :nobj, :nv, :] = s['ref_obj_positions']
                if has_rank:
                    if 'vtx_rank' not in s:
                        raise KeyError("Batch sample missing 'vtx_rank' while others provide it")
                    r = s['vtx_rank'].to(device=device, dtype=torch.int64)
                    vtx_rank[b, :nobj, :nv] = torch.where(r[:, :nv] >= nv, torch.full_like(r[:, :nv], int(max_nv)), r[:, :nv])
                if has_anchors:
                    if 'anchor_idx' not in s:
                        raise KeyError("Batch sample missing 'anchor_idx' while others provide it")
                    if s['anchor_idx'].dim() != 2:
                        raise ValueError(f"anchor_idx must be (Nobj, Nq); got {tuple(s['anchor_idx'].shape)}")
                    Nq = int(s['anchor_idx'].shape[1])
                    if Nq != anchor_k:
                        raise ValueError(f"Inconsistent anchor_k across batch: {Nq} vs {anchor_k}")
                    if s['anchor_idx'].shape[0] != nobj:
                        raise ValueError(f"anchor_idx Nobj ({s['anchor_idx'].shape[0]}) != sample Nobj ({nobj})")
                    anchor_idx[b, :nobj, :] = s['anchor_idx'].to(device=device, dtype=torch.int64)
                    # With strict loader checks, all anchors are valid; set mask True
                    anchor_mask[b, :nobj, :] = True

            collated = {
                'obj_positions': obj_positions,
                'valid_obj_mask': valid_obj_mask,
                'valid_vtx_mask': valid_vtx_mask,
                'ref_obj_positions': ref_obj_positions,
                'dt': torch.stack([s['dt'] for s in batch], dim=0)
            }
            if has_anchors:
                collated['anchor_idx'] = anchor_idx
                collated['anchor_mask'] = anchor_mask
            if has_rank:
                collated['vtx_rank'] = vtx_rank

            # Handle phys_params (Nobj, 3) - pad to max_nobj with PADDING_VALUE
            if 'phys_params' in batch[0]:
                phys_params_padded = torch.full((B, max_nobj, 3), PADDING_VALUE, dtype=torch.float32)
                for b, s in enumerate(batch):
                    nobj = s['phys_params'].shape[0]
                    phys_params_padded[b, :nobj, :] = s['phys_params']
                collated['phys_params'] = phys_params_padded

            # Handle skip_step_size (scalar) - can be stacked directly
            if 'skip_step_size' in batch[0]:
                collated['skip_step_size'] = torch.stack([s['skip_step_size'] for s in batch], dim=0)


            # Pass through any remaining keys; tensors are stacked and must agree in shape.
            for key in batch[0].keys():
                if key in collated:
                    continue
                if isinstance(batch[0][key], torch.Tensor):
                    collated[key] = torch.stack([sample[key] for sample in batch], dim=0)
                else:
                    collated[key] = [sample[key] for sample in batch]

            return collated

    return collate_fn


def create_multi_dataloaders(
    data_path: str,
    train_batch_size: int = 32,
    val_batch_size: int = 32,
    eval_batch_size: int = 32,
    sequence_length: int = 10,
    num_workers: int = 4,
    train_ratio: float = 0.9,
    val_ratio: float = 0.1,
    split_file: Optional[str] = 'dataset_split.json',
    distributed: bool = False,
    rank: int = 0,
    world_size: int = 1,
    fast_index: bool = True,
    if_phys_param: bool = False,
    inference_mode: bool = False,
    **kwargs
) -> Tuple[DataLoader, DataLoader, DataLoader]:
    """
    Create train, validation, and test dataloaders from multiple datasets.

    Args:
        data_path: Parent directory containing the dataset folders listed in the split file
        train_batch_size: Batch size for training loader
        val_batch_size: Batch size for validation loader
        eval_batch_size: Batch size for test/evaluation loader (can be larger, no gradients)
        sequence_length: Number of frames per sequence
        num_workers: Number of workers for data loading
        train_ratio: Ratio of datasets for training
        val_ratio: Ratio of datasets for validation
        **kwargs: Additional arguments for dataset

    Returns:
        train_loader, val_loader, test_loader
    """
    # Sampler options (with-replacement support)
    replacement_sampling: bool = bool(kwargs.pop('replacement_sampling', False))
    replacement_steps_per_epoch = kwargs.pop('replacement_steps_per_epoch', None)

    # Extract padding config for collate_fn (keep in kwargs for dataset creation)
    max_objects_per_frame = kwargs.get('max_objects_per_frame', None)
    max_vertices_per_object = kwargs.get('max_vertices_per_object', None)

    # Add if_phys_param to kwargs for dataset creation
    kwargs['if_phys_param'] = if_phys_param

    # Create configured collate function with fixed padding
    collate_fn = make_collate_fn(
        max_objects_per_frame=max_objects_per_frame,
        max_vertices_per_object=max_vertices_per_object
    )

    # Default cap on vertices per object to bound memory usage.
    if 'max_vertices_per_object' not in kwargs:
        kwargs['max_vertices_per_object'] = 1024

    # Inference mode skips train dataset indexing entirely.
    if inference_mode:
        logger.info("Inference mode: Skipping train dataset indexing")
        train_dataset = None
    else:
        train_dataset = MultiObjectDynamics(
            data_path,
            sequence_length=sequence_length,
            split='train',
            train_ratio=train_ratio,
            val_ratio=val_ratio,
            split_file=split_file,
            fast_index=fast_index,
            **kwargs
        )

    # Only create val/test datasets if validation is enabled (val_ratio > 0) OR inference mode
    if val_ratio > 0.0 or inference_mode:
        val_dataset = MultiObjectDynamics(
            data_path,
            sequence_length=sequence_length,
            split='val',
            train_ratio=train_ratio,
            val_ratio=val_ratio,
            split_file=split_file,
            fast_index=fast_index,
            **kwargs
        )

        test_dataset = MultiObjectDynamics(
            data_path,
            sequence_length=sequence_length,
            split='test',
            train_ratio=train_ratio,
            val_ratio=val_ratio,
            split_file=split_file,
            fast_index=fast_index,
            **kwargs
        )
    else:
        # Validation disabled: create minimal dummy datasets to avoid errors
        logger.info("Validation disabled (val_ratio=0.0) - skipping val/test dataset creation")
        val_dataset = None
        test_dataset = None
    
    # Create distributed samplers when required
    train_sampler = None
    val_sampler = None
    test_sampler = None
    if distributed and train_dataset is not None:
        # For small datasets (overfit mode), disable drop_last to avoid empty loaders
        # drop_last=True only when dataset is large enough (>= world_size * train_batch_size)
        train_drop_last = len(train_dataset) >= world_size * train_batch_size
        if replacement_sampling:
            # Replacement-based distributed sampler
            class ReplacementDistributedSampler(Sampler):
                def __init__(self, dataset, num_replicas, rank, num_samples_per_rank, seed: int = 0):
                    self.dataset = dataset
                    self.num_replicas = num_replicas
                    self.rank = rank
                    self.num_samples = int(num_samples_per_rank)
                    self.seed = int(seed)
                    self.epoch = 0

                def __iter__(self):
                    g = torch.Generator()
                    g.manual_seed(self.seed + self.epoch * 997 + self.rank * 23)
                    total = len(self.dataset)
                    # Draw with replacement
                    indices = torch.randint(high=total, size=(self.num_samples,), generator=g, dtype=torch.int64)
                    return iter(indices.tolist())

                def __len__(self):
                    return self.num_samples

                def set_epoch(self, epoch: int):
                    self.epoch = int(epoch)

            # Determine per-rank target samples
            if replacement_steps_per_epoch is None:
                # Heuristic: ceil per-rank steps from dataset size
                steps = max(1, (len(train_dataset) + world_size * train_batch_size - 1) // (world_size * train_batch_size))
            else:
                steps = int(replacement_steps_per_epoch)
            per_rank_samples = steps * train_batch_size
            train_sampler = ReplacementDistributedSampler(
                train_dataset, num_replicas=world_size, rank=rank, num_samples_per_rank=per_rank_samples, seed=42
            )
        else:
            train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=train_drop_last)

    # Create val/test samplers for distributed mode (separate from train_dataset condition)
    if distributed:
        if val_dataset is not None:
            val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False)
        if test_dataset is not None:
            test_sampler = DistributedSampler(test_dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False)

    # Create dataloaders
    # For DataLoader drop_last: use same logic as DistributedSampler
    if train_dataset is not None:
        dataloader_drop_last = len(train_dataset) >= (world_size if distributed else 1) * train_batch_size
    else:
        dataloader_drop_last = False

    # Handle None or empty train dataset (inference mode)
    if train_dataset is None:
        logger.info("Inference mode: Creating empty train DataLoader (no training data)")
        # Create a minimal empty dataset for the DataLoader
        class EmptyDataset(torch.utils.data.Dataset):
            def __len__(self):
                return 0
            def __getitem__(self, idx):
                raise IndexError("Empty dataset")
        train_loader = DataLoader(
            EmptyDataset(),
            batch_size=1,
            shuffle=False,
            num_workers=0,
            collate_fn=collate_fn
        )
    elif distributed:
        train_loader = DataLoader(
            train_dataset,
            batch_size=train_batch_size,
            shuffle=False,
            sampler=train_sampler,
            num_workers=num_workers,
            collate_fn=collate_fn,
            pin_memory=True,
            persistent_workers=True if num_workers and num_workers > 0 else False,
            prefetch_factor=2 if num_workers and num_workers > 0 else None,
            drop_last=dataloader_drop_last
        )
    else:
        # Single GPU: optional with-replacement RandomSampler
        if replacement_sampling:
            if replacement_steps_per_epoch is None:
                steps = max(1, (len(train_dataset) + train_batch_size - 1) // train_batch_size)
            else:
                steps = int(replacement_steps_per_epoch)
            num_samples = steps * train_batch_size
            single_sampler = RandomSampler(train_dataset, replacement=True, num_samples=num_samples)
            train_loader = DataLoader(
                train_dataset,
                batch_size=train_batch_size,
                shuffle=False,
                sampler=single_sampler,
                num_workers=num_workers,
                collate_fn=collate_fn,
                pin_memory=True,
                persistent_workers=True if num_workers and num_workers > 0 else False,
                prefetch_factor=2 if num_workers and num_workers > 0 else None,
                drop_last=False
            )
        else:
            train_loader = DataLoader(
                train_dataset,
                batch_size=train_batch_size,
                shuffle=True,
                sampler=None,
                num_workers=num_workers,
                collate_fn=collate_fn,
                pin_memory=True,
                persistent_workers=True if num_workers and num_workers > 0 else False,
                prefetch_factor=2 if num_workers and num_workers > 0 else None,
                drop_last=dataloader_drop_last
            )

    # Only create val/test loaders if datasets exist
    if val_dataset is not None:
        val_loader = DataLoader(
            val_dataset,
            batch_size=val_batch_size,
            shuffle=False,
            sampler=val_sampler,
            num_workers=num_workers,
            collate_fn=collate_fn,
            pin_memory=True,
            persistent_workers=True if num_workers and num_workers > 0 else False,
            prefetch_factor=2 if num_workers and num_workers > 0 else None,
            drop_last=False
        )
    else:
        val_loader = None

    if test_dataset is not None:
        test_loader = DataLoader(
            test_dataset,
            batch_size=eval_batch_size,
            shuffle=False,
            sampler=test_sampler,
            num_workers=num_workers,
            collate_fn=collate_fn,
            pin_memory=True,
            persistent_workers=True if num_workers and num_workers > 0 else False,
            prefetch_factor=2 if num_workers and num_workers > 0 else None,
            drop_last=False
        )
    else:
        test_loader = None

    return train_loader, val_loader, test_loader