datasets/video/memory_selection.py

"""Dataset-side memory selection for DeMemWM latent samples."""

from __future__ import annotations

from typing import Dict, Mapping, MutableMapping, Tuple

import numpy as np
import torch


SEGMENT_KEYS = ("anchor", "dynamic", "revisit")
_VALID_DYNAMIC_POLICIES = ("recent", "event_triggered", "multiview")
_VALID_MULTIVIEW_SELECTORS = ("fov_greedy", "pose_plucker_fps")
_FOV_NUM_POINTS = 10000
_FOV_RADIUS = 30.0
_FOV_HALF_H = 105.0 / 2.0
_FOV_HALF_V = 75.0 / 2.0
_FOV_COS_HALF_H = float(np.cos(np.deg2rad(_FOV_HALF_H)))
_FOV_COS_HALF_V = float(np.cos(np.deg2rad(_FOV_HALF_V)))
_POSE_DISTANCE_SCALE = 30.0
_ANGLE_DISTANCE_SCALE = 180.0
_DETERMINISTIC_UNIT_BALL_CACHE: Dict[Tuple[int, str, torch.dtype], torch.Tensor] = {}


def cfg_get(cfg, key: str, default=None):
    if cfg is None:
        return default
    if isinstance(cfg, Mapping):
        return cfg.get(key, default)
    return getattr(cfg, key, default)


def _memory_candidate_frames(num_frames, target_positions, cfg, split, min_candidate_frame=0):
    target_start = int(target_positions[0])
    causal = split != "training" or bool(cfg_get(cfg, "causal", True))
    stop = target_start if causal else num_frames
    return np.arange(min_candidate_frame, stop, dtype=np.int64)


def _exclude_local_context(candidates: np.ndarray, target_positions: np.ndarray, cfg) -> np.ndarray:
    exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
    if exclusion <= 0 or len(candidates) == 0 or len(target_positions) == 0:
        return candidates

    target_start = int(target_positions[0])
    target_stop = int(target_positions[-1]) + 1
    local_start = max(0, target_start - exclusion)
    return candidates[(candidates < local_start) | (candidates >= target_stop)]


def _exclude_revisit_local_context(candidates: np.ndarray, target_positions: np.ndarray, cfg) -> np.ndarray:
    return _exclude_local_context(candidates, target_positions, cfg)


def _as_pose_array(poses) -> np.ndarray:
    poses = np.asarray(poses, dtype=np.float32)
    if poses.ndim != 2 or poses.shape[1] < 5:
        raise ValueError("poses must have shape [num_frames, >=5]")
    return poses[:, :5]


def _as_pose_tensor(poses) -> torch.Tensor:
    pose_tensor = poses if torch.is_tensor(poses) else torch.as_tensor(poses, dtype=torch.float32)
    pose_tensor = pose_tensor.to(dtype=torch.float32)
    if pose_tensor.ndim != 2 or pose_tensor.shape[1] < 5:
        raise ValueError("poses must have shape [num_frames, >=5]")
    return pose_tensor[:, :5]


def _wrap_degrees(diff: torch.Tensor) -> torch.Tensor:
    return torch.abs(torch.remainder(diff + 180.0, 360.0) - 180.0)


def _pose_directions(poses: torch.Tensor) -> torch.Tensor:
    poses = _as_pose_tensor(poses)
    pitch = torch.deg2rad(poses[:, 3])
    yaw = torch.deg2rad(poses[:, 4])
    cos_pitch = torch.cos(pitch)
    directions = torch.stack(
        [
            torch.sin(yaw) * cos_pitch,
            torch.sin(pitch),
            torch.cos(yaw) * cos_pitch,
        ],
        dim=-1,
    )
    return directions / torch.linalg.vector_norm(directions, dim=-1, keepdim=True).clamp_min(1e-6)


def _deterministic_points_in_unit_ball(num_points: int, device, dtype) -> torch.Tensor:
    """Return deterministic points in the unit ball."""

    num_points = int(num_points)
    device = torch.device(device)
    if num_points <= 0:
        return torch.zeros((0, 3), device=device, dtype=dtype)

    cache_key = (num_points, str(device), dtype)
    cached = _DETERMINISTIC_UNIT_BALL_CACHE.get(cache_key)
    if cached is not None:
        return cached

    compute_dtype = torch.float64 if dtype == torch.float64 else torch.float32
    idx = torch.arange(num_points, device=device, dtype=compute_dtype)
    inv_n = 1.0 / float(num_points)
    z = 1.0 - 2.0 * ((idx + 0.5) * inv_n)
    xy = torch.sqrt((1.0 - z * z).clamp_min(0.0))
    theta = idx * 2.399963229728653
    directions = torch.stack(
        [
            xy * torch.cos(theta),
            xy * torch.sin(theta),
            z,
        ],
        dim=-1,
    )
    radii = torch.pow((idx + 0.5) * inv_n, 1.0 / 3.0)
    points = (directions * radii[:, None]).to(dtype=dtype)

    if len(_DETERMINISTIC_UNIT_BALL_CACHE) >= 8:
        _DETERMINISTIC_UNIT_BALL_CACHE.clear()
    _DETERMINISTIC_UNIT_BALL_CACHE[cache_key] = points
    return points


def _sample_points_in_sphere(center: torch.Tensor) -> torch.Tensor:
    device = center.device
    dtype = center.dtype

    offsets = _deterministic_points_in_unit_ball(_FOV_NUM_POINTS, device, dtype) * _FOV_RADIUS
    return center[None, :3] + offsets


def _inside_fov_points(points: torch.Tensor, poses: torch.Tensor) -> torch.Tensor:
    points = points.to(dtype=torch.float32)
    poses = _as_pose_tensor(poses).to(device=points.device, dtype=torch.float32)
    if points.numel() == 0 or poses.numel() == 0:
        return torch.zeros((poses.shape[0], points.shape[0]), device=points.device, dtype=torch.bool)

    vectors = points[None, :, :] - poses[:, None, :3]
    x = vectors[..., 0]
    y = vectors[..., 1]
    z = vectors[..., 2]
    horizontal = torch.sqrt((x * x + z * z).clamp_min(0.0))
    ray_norm = torch.sqrt((horizontal * horizontal + y * y).clamp_min(0.0))

    yaw = torch.deg2rad(poses[:, 4])[:, None]
    pitch = torch.deg2rad(poses[:, 3])[:, None]
    yaw_aligned = x * torch.sin(yaw) + z * torch.cos(yaw)
    pitch_aligned = horizontal * torch.cos(pitch) + y * torch.sin(pitch)

    # Same yaw/elevation half-angle convention as the previous atan2 path,
    # but without per-point inverse trig or degree wrapping.
    return (yaw_aligned > horizontal * _FOV_COS_HALF_H) & (pitch_aligned > ray_norm * _FOV_COS_HALF_V)


def _target_fov_points(target_poses: torch.Tensor) -> torch.Tensor:
    """Sample a fixed-radius sphere and keep points visible in the target FOV."""

    target_poses = _as_pose_tensor(target_poses)
    if target_poses.shape[0] == 0:
        return torch.zeros((0, 3), device=target_poses.device, dtype=target_poses.dtype)

    points = _sample_points_in_sphere(target_poses[0])
    target_inside = _inside_fov_points(points, target_poses).any(dim=0)
    return points[target_inside]


def _pose_preselect(
    candidates: np.ndarray,
    poses_t: torch.Tensor,
    target_positions: np.ndarray,
    cfg=None,
    *,
    extra_topk: int = 0,
) -> np.ndarray:
    topk = cfg_get(cfg, "pose_preselect_topk", 32)
    if topk is None:
        return candidates
    topk = int(topk)
    if topk > 0:
        topk += max(0, int(extra_topk))
    if topk <= 0 or len(candidates) <= topk:
        return candidates

    candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long)
    targets_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
    candidate_poses = poses_t.index_select(0, candidates_t)
    target_poses = poses_t.index_select(0, targets_t)
    candidate_forward = _pose_directions(candidate_poses)
    target_forward = _pose_directions(target_poses)
    translation_norm = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1) / _FOV_RADIUS
    dot = (candidate_forward[:, None, :] * target_forward[None, :, :]).sum(dim=-1).clamp(-1.0, 1.0)
    angular = torch.acos(dot) / torch.pi
    pose_distance = (translation_norm + angular).min(dim=1).values
    rank = pose_distance.to(dtype=torch.float64) - candidates_t.to(dtype=torch.float64) * 1e-12
    selected = torch.topk(rank, k=min(topk, int(candidates_t.numel())), largest=False, sorted=True).indices
    return candidates_t.index_select(0, selected).cpu().numpy()


def _candidate_fov_masks(poses_t: torch.Tensor, candidates: np.ndarray, target_positions: np.ndarray, cfg=None) -> Tuple[torch.Tensor, torch.Tensor]:
    targets_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
    points = _target_fov_points(poses_t.index_select(0, targets_t))
    if len(candidates) == 0 or points.shape[0] == 0:
        return (
            torch.zeros((len(candidates), 0), device=poses_t.device, dtype=torch.bool),
            torch.zeros((len(candidates),), device=poses_t.device, dtype=torch.float32),
        )

    candidates_t = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long)
    chunk_size = int(cfg_get(cfg, "candidate_chunk_size", 0))
    chunk_size = len(candidates) if chunk_size <= 0 else chunk_size
    inside_parts = []
    score_parts = []
    for start in range(0, len(candidates), chunk_size):
        chunk = candidates_t[start : start + chunk_size]
        inside = _inside_fov_points(points, poses_t.index_select(0, chunk))
        inside_parts.append(inside)
        score_parts.append(inside.float().mean(dim=1))
    return torch.cat(inside_parts, dim=0), torch.cat(score_parts, dim=0)


def _plucker_scores_tensor(candidate_poses: torch.Tensor, target_poses: torch.Tensor, cfg=None) -> torch.Tensor:
    candidate_poses = _as_pose_tensor(candidate_poses)
    target_poses = _as_pose_tensor(target_poses).to(device=candidate_poses.device)
    if candidate_poses.shape[0] == 0 or target_poses.shape[0] == 0:
        return torch.zeros((candidate_poses.shape[0],), device=candidate_poses.device, dtype=torch.float32)

    moment_radius = float(cfg_get(cfg, "plucker_moment_radius", 30.0))
    cand_dir = _pose_directions(candidate_poses)
    target_dir = _pose_directions(target_poses)
    cand_moment = torch.linalg.cross(candidate_poses[:, :3], cand_dir, dim=-1)
    target_moment = torch.linalg.cross(target_poses[:, :3], target_dir, dim=-1)

    direction_sim = (cand_dir @ target_dir.T).clamp(0.0, 1.0)
    moment_dist = torch.linalg.vector_norm(cand_moment[:, None, :] - target_moment[None, :, :], dim=-1)
    moment_sim = torch.exp(-moment_dist / max(moment_radius, 1e-6))
    return (direction_sim * moment_sim).max(dim=1).values.to(dtype=torch.float32)


def _empty_fov_candidate_pool(device) -> dict[str, torch.Tensor]:
    return {
        "candidates_t": torch.empty((0,), device=device, dtype=torch.long),
        "candidate_poses": torch.empty((0, 5), device=device, dtype=torch.float32),
        "inside": torch.zeros((0, 0), device=device, dtype=torch.bool),
        "fov_values": torch.empty((0,), device=device, dtype=torch.float32),
        "plucker": torch.empty((0,), device=device, dtype=torch.float32),
        "gaps": torch.empty((0,), device=device, dtype=torch.long),
        "positive_fov": torch.empty((0,), device=device, dtype=torch.bool),
    }


def _build_fov_candidate_pool(
    poses: np.ndarray,
    candidates: np.ndarray,
    target_positions: np.ndarray,
    cfg,
    *,
    use_plucker: bool,
    preselect: bool = True,
) -> dict[str, torch.Tensor]:
    """Build deterministic FOV/Plucker/coverage data for candidate frames."""

    poses_t = _as_pose_tensor(poses)
    device = poses_t.device
    candidates = np.asarray(candidates, dtype=np.int64)
    target_positions = np.asarray(target_positions, dtype=np.int64)
    if len(candidates) == 0 or len(target_positions) == 0:
        return _empty_fov_candidate_pool(device)

    if preselect:
        candidates = _pose_preselect(candidates, poses_t, target_positions, cfg)
    inside, fov_values = _candidate_fov_masks(poses_t, candidates, target_positions, cfg)
    candidates_t = torch.as_tensor(candidates, device=device, dtype=torch.long)
    candidate_poses = poses_t.index_select(0, candidates_t)
    target_t = torch.as_tensor(target_positions, device=device, dtype=torch.long)
    if use_plucker:
        plucker = _plucker_scores_tensor(candidate_poses, poses_t.index_select(0, target_t), cfg)
    else:
        plucker = torch.zeros((candidates_t.shape[0],), device=device, dtype=torch.float32)

    first_target = torch.full_like(candidates_t, int(target_positions[0]))
    return {
        "candidates_t": candidates_t,
        "candidate_poses": candidate_poses,
        "inside": inside,
        "fov_values": fov_values,
        "plucker": plucker,
        "gaps": first_target - candidates_t,
        "positive_fov": inside.any(dim=1),
    }


def _select_fov_pool_rows(pool: dict[str, torch.Tensor], rows: torch.Tensor) -> dict[str, torch.Tensor]:
    candidates_t = pool["candidates_t"]
    return {
        key: value.index_select(0, rows)
        if torch.is_tensor(value) and value.shape[:1] == candidates_t.shape[:1]
        else value
        for key, value in pool.items()
    }


def _unique_ordered_frames(*arrays: np.ndarray | None) -> np.ndarray:
    values = []
    seen = set()
    for array in arrays:
        if array is None:
            continue
        for value in np.asarray(array, dtype=np.int64):
            frame = int(value)
            if frame not in seen:
                seen.add(frame)
                values.append(frame)
    return np.asarray(values, dtype=np.int64)


def _filter_fov_candidate_pool(
    pool: dict[str, torch.Tensor],
    candidates: np.ndarray,
    *,
    extra_frames: np.ndarray | None = None,
) -> dict[str, torch.Tensor]:
    candidates_t = pool["candidates_t"]
    allowed = _unique_ordered_frames(candidates, extra_frames)
    if len(allowed) == 0:
        rows = torch.empty((0,), device=candidates_t.device, dtype=torch.long)
        return _select_fov_pool_rows(pool, rows)

    if candidates_t.numel() == 0:
        rows = torch.empty((0,), device=candidates_t.device, dtype=torch.long)
        return _select_fov_pool_rows(pool, rows)

    allowed_t = torch.as_tensor(allowed, device=candidates_t.device, dtype=torch.long)
    matches = allowed_t[:, None] == candidates_t[None, :]
    keep_allowed = matches.any(dim=1)
    rows = matches.to(dtype=torch.long).argmax(dim=1).index_select(0, torch.nonzero(keep_allowed, as_tuple=False).flatten())
    if rows.numel() == candidates_t.numel() and bool(torch.equal(rows, torch.arange(candidates_t.numel(), device=candidates_t.device))):
        return pool
    return _select_fov_pool_rows(pool, rows)


def _build_shared_fov_candidate_pool(
    poses: np.ndarray,
    target_positions: np.ndarray,
    cfg,
    split: str,
    *,
    min_candidate_frame: int = 0,
    dynamic_count: int = 0,
    revisit_count: int = 0,
    revisit_excluded: np.ndarray | None = None,
) -> dict[str, torch.Tensor] | None:
    """Build one exact FOV pool over the candidates needed by both selectors."""

    target_positions = np.asarray(target_positions, dtype=np.int64)
    if len(target_positions) == 0:
        return None

    base_candidates = _memory_candidate_frames(
        len(poses),
        target_positions,
        cfg,
        split,
        min_candidate_frame=min_candidate_frame,
    )
    if len(base_candidates) == 0:
        return None

    poses_t = _as_pose_tensor(poses)
    dynamic_candidates = None
    if int(dynamic_count) > 0:
        dynamic_base_candidates = _exclude_local_context(base_candidates, target_positions, cfg)
        # Revisit frames are excluded after selection; include replacement rows
        # that can enter dynamic's pose-preselected top-k after that exclusion.
        dynamic_candidates = _pose_preselect(
            dynamic_base_candidates,
            poses_t,
            target_positions,
            cfg,
            extra_topk=max(0, int(revisit_count)),
        )

    revisit_candidates = None
    if split != "training" and int(revisit_count) > 0:
        revisit_candidates = _exclude_revisit_local_context(base_candidates, target_positions, cfg)
        if revisit_excluded is not None and len(revisit_excluded) > 0 and len(revisit_candidates) > 0:
            revisit_candidates = revisit_candidates[~np.isin(revisit_candidates, np.asarray(revisit_excluded, dtype=np.int64))]
        revisit_candidates = _pose_preselect(revisit_candidates, poses_t, target_positions, cfg)

    pool_candidates = _unique_ordered_frames(dynamic_candidates, revisit_candidates)
    if len(pool_candidates) == 0:
        return None
    return _build_fov_candidate_pool(
        poses,
        pool_candidates,
        target_positions,
        cfg,
        use_plucker=True,
        preselect=False,
    )


def _empty_selection(counts: Mapping[str, int]) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
    indices = {}
    masks = {}
    for key in SEGMENT_KEYS:
        count = int(counts.get(key, 0))
        indices[key] = np.full((count,), -1, dtype=np.int64)
        masks[key] = np.zeros((count,), dtype=bool)
    return indices, masks


def _write_segment(indices: MutableMapping[str, np.ndarray], masks: MutableMapping[str, np.ndarray], key: str, selected) -> None:
    selected = np.asarray(selected, dtype=np.int64)
    count = len(indices[key])
    if count == 0 or len(selected) == 0:
        return
    selected = selected[:count]
    indices[key][: len(selected)] = selected
    masks[key][: len(selected)] = True


def _rng_choice(rng, values: np.ndarray, size: int, *, replace: bool) -> np.ndarray:
    chooser = np.random if rng is None else rng
    return np.asarray(chooser.choice(values, size=int(size), replace=replace), dtype=np.int64)


def _rng_random(rng) -> float:
    return float(np.random.random() if rng is None else rng.random())


def _sample_random_selection(candidates: np.ndarray, count: int, rng=None) -> np.ndarray:
    count = int(count)
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    candidates = np.asarray(candidates, dtype=np.int64)
    candidates = candidates[candidates >= 0]
    if len(candidates) == 0:
        return np.empty((0,), dtype=np.int64)
    return _rng_choice(rng, candidates, count, replace=True)


def _best_row(rows: torch.Tensor, gains: torch.Tensor, fov_values: torch.Tensor, plucker: torch.Tensor, gaps: torch.Tensor, candidates_t: torch.Tensor) -> int:
    active = rows
    for values in (gains, fov_values, plucker, -gaps.to(dtype=torch.float32), -candidates_t.to(dtype=torch.float32)):
        vals = values.index_select(0, active)
        active = active.index_select(0, torch.nonzero(vals == vals.max(), as_tuple=False).flatten())
        if active.numel() == 1:
            break
    return int(active[0].item())


def _deterministic_score_order(scores: torch.Tensor, candidates_t: torch.Tensor) -> torch.Tensor:
    order = np.lexsort((candidates_t.detach().cpu().numpy(), -scores.detach().cpu().numpy()))
    return torch.as_tensor(order, device=candidates_t.device, dtype=torch.long)


def _pose_distance_matrix(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
    if left.shape[0] == 0 or right.shape[0] == 0:
        return torch.zeros((left.shape[0], right.shape[0]), device=left.device, dtype=torch.float32)
    right = right.to(device=left.device, dtype=left.dtype)
    translation = torch.linalg.vector_norm(left[:, None, :3] - right[None, :, :3], dim=-1) / _POSE_DISTANCE_SCALE
    angular = torch.linalg.vector_norm(_wrap_degrees(left[:, None, 3:5] - right[None, :, 3:5]), dim=-1) / _ANGLE_DISTANCE_SCALE
    return translation + angular


def _rank_pose_plucker_candidates(
    poses: np.ndarray,
    candidates: np.ndarray,
    target_positions: np.ndarray,
    cfg,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Return candidate ids, candidate poses, and deterministic relevance scores."""

    poses_t = _as_pose_tensor(poses)
    candidates_t = torch.as_tensor(np.asarray(candidates, dtype=np.int64), device=poses_t.device, dtype=torch.long)
    targets_t = torch.as_tensor(np.asarray(target_positions, dtype=np.int64), device=poses_t.device, dtype=torch.long)
    if candidates_t.numel() == 0 or targets_t.numel() == 0:
        empty_ids = torch.empty((0,), device=poses_t.device, dtype=torch.long)
        empty_poses = torch.empty((0, 5), device=poses_t.device, dtype=torch.float32)
        return empty_ids, empty_poses, torch.empty((0,), device=poses_t.device)

    candidate_poses = poses_t.index_select(0, candidates_t)
    target_poses = poses_t.index_select(0, targets_t)
    spatial_scale = max(float(cfg_get(cfg, "pose_similarity_radius", _FOV_RADIUS)), 1e-6)
    angle_scale = max(float(cfg_get(cfg, "pose_similarity_angle_scale", 180.0)), 1e-6)
    spatial_delta = torch.linalg.vector_norm(candidate_poses[:, None, :3] - target_poses[None, :, :3], dim=-1)
    angle_delta = _wrap_degrees(candidate_poses[:, None, 3:] - target_poses[None, :, 3:]).mean(dim=-1)
    spatial_sim = (1.0 - spatial_delta / spatial_scale).clamp(0.0, 1.0)
    angle_sim = (1.0 - angle_delta / angle_scale).clamp(0.0, 1.0)
    pose_scores = ((spatial_sim + angle_sim) * 0.5).max(dim=1).values

    threshold = float(cfg_get(cfg, "pose_similarity_threshold", 0.6))
    valid = pose_scores >= threshold
    if not bool(valid.any()):
        empty_ids = torch.empty((0,), device=poses_t.device, dtype=torch.long)
        empty_poses = torch.empty((0, 5), device=poses_t.device, dtype=torch.float32)
        return empty_ids, empty_poses, torch.empty((0,), device=poses_t.device)

    valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
    valid_candidates = candidates_t.index_select(0, valid_idx)
    valid_poses = candidate_poses.index_select(0, valid_idx)
    rank_scores = pose_scores.index_select(0, valid_idx)
    if bool(cfg_get(cfg, "training_use_plucker", True)):
        plucker = _plucker_scores_tensor(valid_poses, target_poses, cfg)
        rank_scores = rank_scores + float(cfg_get(cfg, "training_plucker_weight", 1.0)) * plucker

    order = _deterministic_score_order(rank_scores, valid_candidates)
    return valid_candidates.index_select(0, order), valid_poses.index_select(0, order), rank_scores.index_select(0, order)


def _select_by_pose_similarity(
    poses: np.ndarray,
    candidates: np.ndarray,
    target_positions: np.ndarray,
    cfg,
    count: int,
    rng=None,
) -> np.ndarray:
    if count <= 0 or len(candidates) == 0 or len(target_positions) == 0:
        return np.empty((0,), dtype=np.int64)

    ranked_candidates, _, _ = _rank_pose_plucker_candidates(poses, candidates, target_positions, cfg)
    if ranked_candidates.numel() == 0:
        return np.empty((0,), dtype=np.int64)

    selected = ranked_candidates[:count].cpu().numpy()
    return np.sort(np.asarray(selected, dtype=np.int64))


def _select_pose_fps(
    candidate_ids: torch.Tensor,
    candidate_poses: torch.Tensor,
    count: int,
    *,
    seed_ids: np.ndarray | None = None,
    all_poses: np.ndarray | torch.Tensor | None = None,
) -> np.ndarray:
    """Select diverse camera poses from a relevance-filtered candidate set."""

    candidate_ids = torch.as_tensor(candidate_ids, dtype=torch.long)
    candidate_poses = _as_pose_tensor(candidate_poses).to(device=candidate_ids.device)
    if count <= 0 or candidate_ids.numel() == 0:
        return np.empty((0,), dtype=np.int64)

    available = torch.ones((candidate_ids.shape[0],), device=candidate_ids.device, dtype=torch.bool)
    selected_rows = []
    seed_poses = torch.empty((0, 5), device=candidate_ids.device, dtype=torch.float32)
    if seed_ids is not None and all_poses is not None:
        all_poses_t = _as_pose_tensor(all_poses).to(device=candidate_ids.device)
        seed_t = torch.as_tensor(np.asarray(seed_ids, dtype=np.int64), device=candidate_ids.device, dtype=torch.long)
        seed_t = seed_t[(seed_t >= 0) & (seed_t < all_poses_t.shape[0])]
        if seed_t.numel() > 0:
            seed_poses = all_poses_t.index_select(0, torch.unique(seed_t, sorted=True))

    min_dist = None
    if seed_poses.shape[0] > 0:
        min_dist = _pose_distance_matrix(candidate_poses, seed_poses).min(dim=1).values

    for step in range(min(count, int(candidate_ids.numel()))):
        rows = torch.nonzero(available, as_tuple=False).flatten()
        if rows.numel() == 0:
            break
        if step == 0 and min_dist is None:
            row = int(rows[0].item())
        else:
            vals = min_dist.index_select(0, rows)
            tied = rows.index_select(0, torch.nonzero(vals == vals.max(), as_tuple=False).flatten())
            row = int(tied[torch.argmin(candidate_ids.index_select(0, tied))].item())
        selected_rows.append(row)
        available[row] = False

        dists = _pose_distance_matrix(candidate_poses, candidate_poses[row : row + 1]).flatten()
        min_dist = dists if min_dist is None else torch.minimum(min_dist, dists)
        min_dist = min_dist.masked_fill(~available, float("-inf"))

    selected = candidate_ids.index_select(0, torch.as_tensor(selected_rows, device=candidate_ids.device, dtype=torch.long))
    return selected.cpu().numpy().astype(np.int64)


def _select_dynamic_multiview_pose_plucker_fps(
    poses: np.ndarray,
    candidates: np.ndarray,
    target_positions: np.ndarray,
    cfg,
    count: int,
    *,
    seed_ids: np.ndarray | None = None,
) -> np.ndarray:
    ranked_ids, ranked_poses, _ = _rank_pose_plucker_candidates(poses, candidates, target_positions, cfg)
    topk = cfg_get(cfg, "pose_preselect_topk", 32)
    if topk is not None:
        topk = int(topk)
        if topk > 0:
            ranked_ids = ranked_ids[:topk]
            ranked_poses = ranked_poses[:topk]
    return _select_pose_fps(ranked_ids, ranked_poses, count, seed_ids=seed_ids, all_poses=poses)


def _select_dynamic_multiview_from_fov_pool(
    pool: dict[str, torch.Tensor],
    count: int,
    *,
    excluded: np.ndarray | None = None,
    reference_frames: np.ndarray | None = None,
    all_poses: np.ndarray | torch.Tensor | None = None,
) -> np.ndarray:
    """Select dynamic multiview frames by residual FOV coverage."""

    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    candidates_t = pool["candidates_t"]
    candidate_poses = pool["candidate_poses"]
    if candidates_t.numel() == 0:
        return np.empty((0,), dtype=np.int64)

    valid = torch.ones((candidates_t.shape[0],), device=candidates_t.device, dtype=torch.bool)
    if excluded is not None and len(excluded) > 0:
        excluded_t = torch.as_tensor(np.asarray(excluded, dtype=np.int64), device=candidates_t.device, dtype=torch.long)
        valid &= ~(candidates_t[:, None] == excluded_t[None, :]).any(dim=1)
    if not bool(valid.any()):
        return np.empty((0,), dtype=np.int64)

    positive = pool.get("positive_fov", torch.zeros_like(valid)).to(device=candidates_t.device, dtype=torch.bool) & valid
    if not bool(positive.any()):
        rows = torch.nonzero(valid, as_tuple=False).flatten()
        return _select_pose_fps(
            candidates_t.index_select(0, rows),
            candidate_poses.index_select(0, rows),
            count,
            seed_ids=reference_frames,
            all_poses=all_poses,
        )

    rows = torch.nonzero(positive, as_tuple=False).flatten()
    candidates_t = candidates_t.index_select(0, rows)
    candidate_poses = candidate_poses.index_select(0, rows)
    inside = pool["inside"].index_select(0, rows)
    fov_values = pool["fov_values"].index_select(0, rows)
    plucker = pool["plucker"].index_select(0, rows)
    gaps = pool["gaps"].index_select(0, rows)
    temporal_gap = gaps.abs()

    covered = torch.zeros((inside.shape[1],), device=candidates_t.device, dtype=torch.bool)
    reference_poses = []
    if reference_frames is not None and len(reference_frames) > 0:
        ref_t = torch.as_tensor(np.asarray(reference_frames, dtype=np.int64), device=candidates_t.device, dtype=torch.long)
        ref_rows = torch.nonzero((pool["candidates_t"][:, None] == ref_t[None, :]).any(dim=1), as_tuple=False).flatten()
        if ref_rows.numel() > 0:
            covered |= pool["inside"].index_select(0, ref_rows).any(dim=0)
            reference_poses.append(pool["candidate_poses"].index_select(0, ref_rows))
        if all_poses is not None:
            all_poses_t = _as_pose_tensor(all_poses).to(device=candidates_t.device)
            ref_t = ref_t[(ref_t >= 0) & (ref_t < all_poses_t.shape[0])]
            if ref_t.numel() > 0:
                reference_poses.append(all_poses_t.index_select(0, torch.unique(ref_t, sorted=True)))

    selected_rows = []
    remaining = torch.ones((candidates_t.shape[0],), device=candidates_t.device, dtype=torch.bool)
    for _ in range(min(count, int(candidates_t.numel()))):
        active = torch.nonzero(remaining, as_tuple=False).flatten()
        if active.numel() == 0:
            break

        if selected_rows:
            selected_t = torch.as_tensor(selected_rows, device=candidates_t.device, dtype=torch.long)
            refs = reference_poses + [candidate_poses.index_select(0, selected_t)]
        else:
            refs = reference_poses
        if refs:
            reference_pose_t = torch.cat(refs, dim=0)
            pose_dist = _pose_distance_matrix(candidate_poses, reference_pose_t).min(dim=1).values
        else:
            pose_dist = torch.zeros((candidate_poses.shape[0],), device=candidates_t.device, dtype=torch.float32)
        gains = (inside & ~covered[None, :]).float().mean(dim=1)

        row = active
        for values in (
            gains,
            fov_values,
            plucker,
            pose_dist,
            -temporal_gap.to(dtype=torch.float32),
            -candidates_t.to(dtype=torch.float32),
        ):
            vals = values.index_select(0, row)
            row = row.index_select(0, torch.nonzero(vals == vals.max(), as_tuple=False).flatten())
            if row.numel() == 1:
                break
        selected_row = int(row[0].item())
        selected_rows.append(selected_row)
        covered |= inside[selected_row]
        remaining[selected_row] = False

    selected = candidates_t.index_select(0, torch.as_tensor(selected_rows, device=candidates_t.device, dtype=torch.long))
    return selected.cpu().numpy().astype(np.int64)


def _select_by_point_union(
    poses: np.ndarray,
    candidates: np.ndarray,
    target_positions: np.ndarray,
    cfg,
    count: int,
    *,
    fov_threshold: float | None = None,
    min_total_coverage: float | None = None,
    use_plucker: bool = False,
    fov_pool: dict[str, torch.Tensor] | None = None,
) -> np.ndarray:
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    candidates = np.asarray(candidates, dtype=np.int64)
    target_positions = np.asarray(target_positions, dtype=np.int64)
    if len(candidates) == 0 or len(target_positions) == 0:
        return np.empty((0,), dtype=np.int64)

    if fov_pool is None:
        fov_pool = _build_fov_candidate_pool(poses, candidates, target_positions, cfg, use_plucker=use_plucker)
    else:
        candidates = _pose_preselect(candidates, _as_pose_tensor(poses), target_positions, cfg)
        fov_pool = _filter_fov_candidate_pool(fov_pool, candidates)

    return _select_by_point_union_from_pool(
        fov_pool,
        count,
        fov_threshold=fov_threshold,
        min_total_coverage=min_total_coverage,
    )


def _select_by_point_union_from_pool(
    pool: dict[str, torch.Tensor],
    count: int,
    *,
    fov_threshold: float | None = None,
    min_total_coverage: float | None = None,
) -> np.ndarray:
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    candidates_t = pool["candidates_t"]
    inside = pool["inside"]
    if candidates_t.numel() == 0 or inside.shape[0] == 0 or inside.shape[1] == 0:
        return np.empty((0,), dtype=np.int64)

    device = candidates_t.device
    fov_values = pool["fov_values"]
    plucker = pool["plucker"]
    gaps = pool["gaps"]
    valid = torch.ones((candidates_t.shape[0],), device=device, dtype=torch.bool)
    if fov_threshold is not None:
        valid &= fov_values >= float(fov_threshold)
    if not bool(valid.any()):
        return np.empty((0,), dtype=np.int64)

    valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
    candidates_t = candidates_t.index_select(0, valid_idx)
    inside = inside.index_select(0, valid_idx)
    fov_values = fov_values.index_select(0, valid_idx)
    plucker = plucker.index_select(0, valid_idx)
    gaps = gaps.index_select(0, valid_idx)

    remaining = torch.ones((candidates_t.shape[0],), device=device, dtype=torch.bool)
    covered = torch.zeros((inside.shape[1],), device=device, dtype=torch.bool)
    selected_rows = []

    for _ in range(count):
        gains = (inside & ~covered[None, :]).float().mean(dim=1)
        rows = torch.nonzero(remaining, as_tuple=False).flatten()
        if rows.numel() == 0:
            break
        row = _best_row(rows, gains, fov_values, plucker, gaps, candidates_t)
        if float(gains[row].item()) <= 0.0 and float(fov_values[row].item()) <= 0.0:
            break
        selected_rows.append(row)
        covered |= inside[row]
        remaining[row] = False

    if not selected_rows:
        return np.empty((0,), dtype=np.int64)

    coverage = float(covered.float().mean().item()) if covered.numel() else 0.0
    if min_total_coverage is not None and coverage < float(min_total_coverage):
        return np.empty((0,), dtype=np.int64)
    selected = candidates_t.index_select(0, torch.as_tensor(selected_rows, device=device, dtype=torch.long))
    return np.sort(selected.cpu().numpy().astype(np.int64))


def _select_anchor(candidates: np.ndarray, count: int, cfg, poses=None) -> np.ndarray:
    if count <= 0 or len(candidates) == 0:
        return np.empty((0,), dtype=np.int64)
    if not bool(cfg_get(cfg, "anchor_diverse_selection", True)) or len(candidates) <= count or poses is None:
        return candidates[:count].astype(np.int64)

    poses_t = _as_pose_tensor(poses)
    candidates_t = torch.as_tensor(candidates.astype(np.int64), device=poses_t.device, dtype=torch.long)
    candidate_poses = poses_t.index_select(0, candidates_t).float()
    spatial = torch.cdist(candidate_poses[:, :3], candidate_poses[:, :3])
    angular = torch.linalg.vector_norm(
        _wrap_degrees(candidate_poses[:, None, 3:] - candidate_poses[None, :, 3:]),
        dim=-1,
    )
    pairwise = torch.sqrt(spatial.square() + angular.square())
    if not bool((pairwise > 0).any().item()):
        return candidates[:count].astype(np.int64)

    available = torch.ones((int(candidates_t.numel()),), device=poses_t.device, dtype=torch.bool)
    selected = [0]
    available[0] = False

    dists = pairwise[selected].min(dim=0).values.masked_fill(~available, float("-inf"))
    for _ in range(count - len(selected)):
        farthest = int(dists.argmax().item())
        if not bool(available[farthest].item()):
            break
        selected.append(farthest)
        available[farthest] = False
        dists = torch.minimum(dists, pairwise[farthest]).masked_fill(~available, float("-inf"))

    selected_t = torch.as_tensor(sorted(selected[:count]), device=poses_t.device, dtype=torch.long)
    return candidates_t.index_select(0, selected_t).cpu().numpy().astype(np.int64)


def _select_dynamic(target_start: int, count: int, min_candidate_frame: int = 0) -> np.ndarray:
    if count <= 0 or target_start <= min_candidate_frame:
        return np.empty((0,), dtype=np.int64)
    start = max(min_candidate_frame, target_start - count)
    return np.arange(start, target_start, dtype=np.int64)[-count:]


def _dynamic_policy(cfg) -> str:
    dynamic_cfg = cfg_get(cfg, "dynamic", {})
    policy = str(cfg_get(dynamic_cfg, "selection_policy", "recent"))
    if policy not in _VALID_DYNAMIC_POLICIES:
        valid = ", ".join(_VALID_DYNAMIC_POLICIES)
        raise ValueError(f"memory_selection.dynamic.selection_policy must be one of {valid}; got {policy!r}")
    return policy


def _dynamic_multiview_selector(cfg) -> str:
    dynamic_cfg = cfg_get(cfg, "dynamic", {})
    selector = str(cfg_get(dynamic_cfg, "multiview_selector", "fov_greedy"))
    if selector not in _VALID_MULTIVIEW_SELECTORS:
        valid = ", ".join(_VALID_MULTIVIEW_SELECTORS)
        raise ValueError(f"memory_selection.dynamic.multiview_selector must be one of {valid}; got {selector!r}")
    return selector


def _select_dynamic_multiview(
    poses: np.ndarray,
    target_positions: np.ndarray,
    cfg,
    count: int,
    *,
    excluded: np.ndarray | None = None,
    reference_frames: np.ndarray | None = None,
    split: str = "training",
    min_candidate_frame: int = 0,
    fov_pool: dict[str, torch.Tensor] | None = None,
) -> np.ndarray:
    """Select dynamic-as-multiview memory frames."""

    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    poses = _as_pose_array(poses)
    target_positions = np.asarray(target_positions, dtype=np.int64)
    if len(target_positions) == 0:
        return np.empty((0,), dtype=np.int64)

    candidates = _memory_candidate_frames(
        len(poses),
        target_positions,
        cfg,
        split,
        min_candidate_frame=min_candidate_frame,
    )
    candidates = _exclude_local_context(candidates, target_positions, cfg)
    if excluded is not None and len(excluded) > 0 and len(candidates) > 0:
        candidates = candidates[~np.isin(candidates, np.asarray(excluded, dtype=np.int64))]
    if len(candidates) == 0:
        return np.empty((0,), dtype=np.int64)

    selector = _dynamic_multiview_selector(cfg)
    if selector == "fov_greedy":
        pool = fov_pool
        if pool is None:
            pool = _build_fov_candidate_pool(poses, candidates, target_positions, cfg, use_plucker=True)
        else:
            candidates = _pose_preselect(candidates, _as_pose_tensor(poses), target_positions, cfg)
            # Reference rows seed residual coverage below even when excluded from selection.
            pool = _filter_fov_candidate_pool(pool, candidates, extra_frames=reference_frames)
        selected = _select_dynamic_multiview_from_fov_pool(
            pool,
            count,
            excluded=excluded,
            reference_frames=reference_frames,
            all_poses=poses,
        )
    elif selector == "pose_plucker_fps":
        selected = _select_dynamic_multiview_pose_plucker_fps(
            poses,
            candidates,
            target_positions,
            cfg,
            count,
            seed_ids=reference_frames,
        )
    else:
        raise AssertionError(f"unhandled dynamic multiview selector: {selector!r}")

    return np.sort(np.asarray(selected, dtype=np.int64))


def _latent_frame_vectors(latents, start: int, stop: int):
    if latents is None or stop <= start:
        return None
    length = len(latents)
    start = max(0, min(int(start), length))
    stop = max(start, min(int(stop), length))
    if stop <= start:
        return None

    values = latents[start:stop]
    tensor = values if torch.is_tensor(values) else torch.as_tensor(np.array(values), dtype=torch.float32)
    tensor = tensor.to(dtype=torch.float32)
    if tensor.ndim == 1:
        return tensor[:, None]
    if tensor.ndim == 2:
        return tensor
    return tensor.reshape(tensor.shape[0], tensor.shape[1], -1).mean(dim=-1)


def _frame_l2_deltas(values, frames: np.ndarray, history_start: int, stop: int, device) -> torch.Tensor:
    deltas = torch.zeros((len(frames),), device=device, dtype=torch.float32)
    if values is None or len(frames) == 0:
        return deltas

    length = len(values)
    value_start = max(0, min(int(history_start), length))
    value_stop = max(value_start, min(int(stop), length))
    if value_stop - value_start <= 1:
        return deltas

    sliced = values[value_start:value_stop]
    tensor = sliced if torch.is_tensor(sliced) else torch.as_tensor(np.asarray(sliced), dtype=torch.float32, device=device)
    tensor = tensor.to(device=device, dtype=torch.float32)
    if tensor.ndim == 1:
        tensor = tensor[:, None]
    else:
        tensor = tensor.reshape(tensor.shape[0], -1)

    consecutive = torch.linalg.vector_norm(tensor[1:] - tensor[:-1], dim=-1)
    valid = (frames >= value_start + 1) & (frames < value_stop)
    if bool(valid.any()):
        rows = torch.as_tensor(frames[valid] - (value_start + 1), device=device, dtype=torch.long)
        deltas[torch.as_tensor(valid, device=device, dtype=torch.bool)] = consecutive.index_select(0, rows)
    return deltas


def _pose_delta_values(poses, frames: np.ndarray, history_start: int, stop: int, device) -> torch.Tensor:
    deltas = torch.zeros((len(frames),), device=device, dtype=torch.float32)
    if poses is None or len(frames) == 0:
        return deltas

    length = len(poses)
    value_start = max(0, min(int(history_start), length))
    value_stop = max(value_start, min(int(stop), length))
    if value_stop - value_start <= 1:
        return deltas

    pose_tensor = _as_pose_tensor(poses[value_start:value_stop]).to(device=device, dtype=torch.float32)
    # Consecutive pose novelty keeps yaw/pitch in degrees modulo 360.
    spatial = torch.linalg.vector_norm(pose_tensor[1:, :3] - pose_tensor[:-1, :3], dim=-1) / _POSE_DISTANCE_SCALE
    angular = torch.linalg.vector_norm(_wrap_degrees(pose_tensor[1:, 3:5] - pose_tensor[:-1, 3:5]), dim=-1) / _ANGLE_DISTANCE_SCALE
    consecutive = spatial + angular

    valid = (frames >= value_start + 1) & (frames < value_stop)
    if bool(valid.any()):
        rows = torch.as_tensor(frames[valid] - (value_start + 1), device=device, dtype=torch.long)
        deltas[torch.as_tensor(valid, device=device, dtype=torch.bool)] = consecutive.index_select(0, rows)
    return deltas


def _robust_z(values: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    if values.numel() == 0:
        return values
    median = values.median()
    mad = (values - median).abs().median()
    return ((values - median) / (1.4826 * mad + eps)).clamp(0.0, 8.0)


def _event_triggered_anchor_candidates_from_deltas(
    frames: np.ndarray,
    d_vis,
    d_pose,
    d_act,
    cfg,
) -> np.ndarray:
    frames = np.asarray(frames, dtype=np.int64)
    if len(frames) == 0:
        return np.empty((0,), dtype=np.int64)

    d_vis = d_vis if torch.is_tensor(d_vis) else torch.as_tensor(d_vis, dtype=torch.float32)
    device = d_vis.device
    d_vis = d_vis.to(dtype=torch.float32)
    d_pose = torch.as_tensor(d_pose, device=device, dtype=torch.float32)
    d_act = torch.as_tensor(d_act, device=device, dtype=torch.float32)

    dynamic_cfg = cfg_get(cfg, "dynamic", {})
    expected_vis = float(cfg_get(dynamic_cfg, "b_pose", 0.3)) * d_pose + float(cfg_get(dynamic_cfg, "b_action", 0.2)) * d_act
    r_vis = (d_vis - expected_vis).clamp_min(0.0)

    z_vis = _robust_z(d_vis)
    z_pose = _robust_z(d_pose)
    z_act = _robust_z(d_act)
    z_r_vis = _robust_z(r_vis)
    scene_scores = z_r_vis + 0.6 * z_vis + 0.4 * z_pose
    state_scores = 1.2 * z_r_vis + 0.8 * z_act
    quality = 0.5 * z_vis + 0.3 * z_pose + 0.2 * z_act

    scene_threshold = float(cfg_get(dynamic_cfg, "scene_threshold", 2.5))
    state_threshold = float(cfg_get(dynamic_cfg, "state_threshold", 2.5))
    stable_threshold = float(cfg_get(dynamic_cfg, "stable_threshold", 1.0))
    stable_frames = max(1, int(cfg_get(dynamic_cfg, "stable_frames", 3)))
    min_event_gap = max(0, int(cfg_get(dynamic_cfg, "min_event_gap", 8)))
    min_anchor_score = float(cfg_get(dynamic_cfg, "min_anchor_score", 2.0))

    anchors = []
    in_event = False
    peak_score = 0.0
    stable_count = 0
    stable_start = -1
    cooldown_until = -1
    trigger_scores = torch.maximum(scene_scores, state_scores).detach().cpu().numpy()
    scene_np = scene_scores.detach().cpu().numpy()
    state_np = state_scores.detach().cpu().numpy()
    quality_np = quality.detach().cpu().numpy()

    for frame, scene_score, state_score, trigger_score, q_value in zip(frames, scene_np, state_np, trigger_scores, quality_np):
        frame = int(frame)
        if frame < cooldown_until:
            continue

        triggered = scene_score >= scene_threshold or state_score >= state_threshold
        if not in_event:
            if triggered:
                in_event = True
                peak_score = float(trigger_score)
                stable_count = 0
                stable_start = -1
            continue

        peak_score = max(peak_score, float(trigger_score))
        if float(q_value) < stable_threshold:
            if stable_count == 0:
                stable_start = frame
            stable_count += 1
            if stable_count >= stable_frames:
                if peak_score >= min_anchor_score:
                    anchors.append(stable_start)
                in_event = False
                stable_count = 0
                stable_start = -1
                cooldown_until = frame + min_event_gap
        else:
            stable_count = 0
            stable_start = -1

    return np.sort(np.asarray(anchors, dtype=np.int64))


def _event_triggered_anchor_candidates(
    target_start: int,
    cfg,
    poses=None,
    latents=None,
    actions=None,
    min_candidate_frame: int = 0,
) -> np.ndarray:
    if latents is None or target_start <= min_candidate_frame:
        return np.empty((0,), dtype=np.int64)

    history_start = max(0, int(min_candidate_frame) - 1)
    vectors = _latent_frame_vectors(latents, history_start, int(target_start))
    if vectors is None or vectors.shape[0] == 0:
        return np.empty((0,), dtype=np.int64)

    stop = history_start + int(vectors.shape[0])
    frames = np.arange(int(min_candidate_frame), stop, dtype=np.int64)
    if len(frames) == 0:
        return np.empty((0,), dtype=np.int64)

    norms = torch.linalg.vector_norm(vectors, dim=-1, keepdim=True)
    normalized = vectors / norms.clamp_min(1e-6)
    cosine = (normalized[1:] * normalized[:-1]).sum(dim=-1).clamp(-1.0, 1.0)
    valid_pair = (norms[1:, 0] > 1e-6) & (norms[:-1, 0] > 1e-6)
    cosine = torch.where(valid_pair, cosine, torch.ones_like(cosine))
    consecutive_vis = 1.0 - cosine
    d_vis = torch.zeros((len(frames),), device=vectors.device, dtype=torch.float32)
    valid_vis = (frames >= history_start + 1) & (frames < stop)
    if bool(valid_vis.any()):
        rows = torch.as_tensor(frames[valid_vis] - (history_start + 1), device=vectors.device, dtype=torch.long)
        d_vis[torch.as_tensor(valid_vis, device=vectors.device, dtype=torch.bool)] = consecutive_vis.index_select(0, rows)

    d_pose = _pose_delta_values(poses, frames, history_start, stop, vectors.device)
    d_act = _frame_l2_deltas(actions, frames, history_start, stop, vectors.device)
    return _event_triggered_anchor_candidates_from_deltas(frames, d_vis, d_pose, d_act, cfg)


def _nearest_unique_event_anchors(event_anchors: np.ndarray, reference_frames: np.ndarray, count: int) -> np.ndarray:
    if count <= 0 or len(event_anchors) == 0:
        return np.empty((0,), dtype=np.int64)
    references = np.asarray(reference_frames, dtype=np.int64)
    if len(references) == 0:
        return np.empty((0,), dtype=np.int64)

    selected = []
    used = set()
    base_quota, extra = divmod(int(count), len(references))
    for ref_idx, ref in enumerate(references):
        quota = base_quota + (1 if ref_idx < extra else 0)
        if quota <= 0:
            continue
        order = np.lexsort((event_anchors, np.abs(event_anchors - int(ref))))
        picked_for_ref = 0
        for row in order:
            anchor = int(event_anchors[row])
            if anchor in used:
                continue
            selected.append(anchor)
            used.add(anchor)
            picked_for_ref += 1
            if len(selected) >= count or picked_for_ref >= quota:
                break

    if len(selected) < count:
        distance = np.min(np.abs(event_anchors[:, None] - references[None, :]), axis=1)
        order = np.lexsort((event_anchors, distance))
        for row in order:
            anchor = int(event_anchors[row])
            if anchor in used:
                continue
            selected.append(anchor)
            used.add(anchor)
            if len(selected) >= count:
                break

    return np.sort(np.asarray(selected[:count], dtype=np.int64))


def _build_dynamic_stream(
    latents,
    actions=None,
    poses=None,
    cfg=None,
    min_candidate_frame: int = 0,
    stop: int | None = None,
) -> np.ndarray:
    if latents is None:
        return np.empty((0,), dtype=np.int64)

    length = len(latents)
    min_candidate_frame = max(0, int(min_candidate_frame))
    target_stop = length if stop is None else max(0, min(int(stop), length))
    if target_stop <= min_candidate_frame:
        return np.empty((0,), dtype=np.int64)

    event_anchors = _event_triggered_anchor_candidates(
        target_stop,
        cfg,
        poses=poses,
        latents=latents,
        actions=actions,
        min_candidate_frame=min_candidate_frame,
    )
    stream = np.concatenate([np.asarray([min_candidate_frame], dtype=np.int64), event_anchors.astype(np.int64, copy=False)])
    stream = stream[(stream >= min_candidate_frame) & (stream < target_stop)]
    return np.unique(np.sort(stream.astype(np.int64, copy=False)))


def _select_dynamic_from_stream(dynamic_stream, reference_frames, count: int) -> np.ndarray:
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    stream = np.asarray(dynamic_stream, dtype=np.int64)
    stream = np.unique(np.sort(stream[stream >= 0]))
    if len(stream) == 0:
        return np.empty((0,), dtype=np.int64)

    references = np.asarray([] if reference_frames is None else reference_frames, dtype=np.int64)
    references = np.unique(np.sort(references[references >= 0]))
    if len(references) == 0:
        return stream[-count:].astype(np.int64, copy=False)

    selected = []
    used = set()

    def add_candidate(value) -> bool:
        value = int(value)
        if value in used:
            return False
        selected.append(value)
        used.add(value)
        return True

    base_quota, extra = divmod(int(count), len(references))
    for ref_idx, ref in enumerate(references):
        quota = base_quota + (1 if ref_idx < extra else 0)
        if quota <= 0:
            continue
        picked = 0
        before = stream[stream <= int(ref)]
        after = stream[stream >= int(ref)]
        brackets = []
        if len(before) > 0:
            brackets.append(int(before[-1]))
        if len(after) > 0:
            brackets.append(int(after[0]))
        for candidate in brackets:
            picked += int(add_candidate(candidate))
            if picked >= quota:
                break
        if picked >= quota:
            continue

        order = np.lexsort((stream, np.abs(stream - int(ref))))
        for row in order:
            picked += int(add_candidate(stream[row]))
            if picked >= quota:
                break

    if len(selected) < count:
        distance = np.min(np.abs(stream[:, None] - references[None, :]), axis=1)
        order = np.lexsort((stream, distance))
        for row in order:
            add_candidate(stream[row])
            if len(selected) >= count:
                break

    return np.sort(np.asarray(selected[:count], dtype=np.int64))


def _select_event_triggered_dynamic(
    target_start: int,
    count: int,
    cfg,
    poses=None,
    latents=None,
    actions=None,
    min_candidate_frame: int = 0,
    reference_frames=None,
    excluded=None,
) -> np.ndarray:
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    dynamic_cfg = cfg_get(cfg, "dynamic", {})
    max_event_anchors = cfg_get(dynamic_cfg, "max_event_anchors")
    if max_event_anchors is not None:
        count = min(count, max(0, int(max_event_anchors)))
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    event_anchors = _event_triggered_anchor_candidates(
        target_start,
        cfg,
        poses=poses,
        latents=latents,
        actions=actions,
        min_candidate_frame=min_candidate_frame,
    )
    event_anchors = _exclude_local_context(event_anchors, np.asarray([target_start], dtype=np.int64), cfg)
    if excluded is not None and len(excluded) > 0 and len(event_anchors) > 0:
        event_anchors = event_anchors[~np.isin(event_anchors, np.asarray(excluded, dtype=np.int64))]
    if len(event_anchors) == 0:
        return np.empty((0,), dtype=np.int64)

    references = np.asarray([target_start] if reference_frames is None or len(reference_frames) == 0 else reference_frames, dtype=np.int64)
    return _nearest_unique_event_anchors(event_anchors, references, count)


def _select_dynamic_by_policy(
    target_start: int,
    count: int,
    cfg,
    poses=None,
    latents=None,
    actions=None,
    min_candidate_frame: int = 0,
    reference_frames=None,
    excluded=None,
    target_positions=None,
    split: str = "training",
    fov_pool: dict[str, torch.Tensor] | None = None,
) -> np.ndarray:
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    policy = _dynamic_policy(cfg)
    if policy == "recent":
        dynamic_stop = int(target_start)
        if target_positions is not None:
            exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
            dynamic_stop = max(min_candidate_frame, dynamic_stop - exclusion) if exclusion > 0 else dynamic_stop
        return _select_dynamic(dynamic_stop, count, min_candidate_frame)
    if policy == "event_triggered":
        return _select_event_triggered_dynamic(
            target_start,
            count,
            cfg,
            poses=poses,
            latents=latents,
            actions=actions,
            min_candidate_frame=min_candidate_frame,
            reference_frames=reference_frames,
            excluded=excluded,
        )
    if policy == "multiview":
        if poses is None:
            raise ValueError("multiview dynamic selection requires poses")
        if target_positions is None:
            target_positions = np.asarray([target_start], dtype=np.int64)
        return _select_dynamic_multiview(
            poses,
            target_positions,
            cfg,
            count,
            excluded=excluded,
            reference_frames=reference_frames,
            split=split,
            min_candidate_frame=min_candidate_frame,
            fov_pool=fov_pool,
        )
    raise AssertionError(f"unhandled dynamic selection policy: {policy!r}")


def _dynamic_random_candidates(
    num_frames: int,
    target_start: int,
    target_positions: np.ndarray,
    cfg,
    split: str,
    min_candidate_frame: int,
    policy: str,
    *,
    excluded: np.ndarray | None = None,
) -> np.ndarray:
    if policy == "recent":
        stop = int(target_start)
        exclusion = max(0, int(cfg_get(cfg, "local_context_exclusion_frames", 8)))
        if exclusion > 0:
            stop = max(int(min_candidate_frame), stop - exclusion)
        return np.arange(int(min_candidate_frame), stop, dtype=np.int64)

    candidates = _memory_candidate_frames(
        int(num_frames),
        target_positions,
        cfg,
        split,
        min_candidate_frame=int(min_candidate_frame),
    )
    candidates = _exclude_local_context(candidates, target_positions, cfg)
    if policy == "multiview" and excluded is not None and len(excluded) > 0 and len(candidates) > 0:
        candidates = candidates[~np.isin(candidates, np.asarray(excluded, dtype=np.int64))]
    return candidates.astype(np.int64, copy=False)


def _select_revisit(
    poses: np.ndarray,
    target_positions: np.ndarray,
    cfg,
    count: int,
    excluded: np.ndarray,
    split: str,
    rng=None,
    min_candidate_frame: int = 0,
    fov_pool: dict[str, torch.Tensor] | None = None,
) -> np.ndarray:
    if count <= 0:
        return np.empty((0,), dtype=np.int64)

    candidates = _memory_candidate_frames(
        len(poses),
        target_positions,
        cfg,
        split,
        min_candidate_frame=min_candidate_frame,
    )
    if len(candidates) == 0:
        return np.empty((0,), dtype=np.int64)

    candidates = _exclude_revisit_local_context(candidates, target_positions, cfg)
    if len(candidates) == 0:
        return np.empty((0,), dtype=np.int64)

    if len(excluded) > 0:
        candidates = candidates[~np.isin(candidates, excluded)]
    if len(candidates) == 0:
        return np.empty((0,), dtype=np.int64)

    if split == "training":
        return _select_by_pose_similarity(poses, candidates, target_positions, cfg, count, rng=rng)

    return _select_by_point_union(
        poses,
        candidates,
        target_positions,
        cfg,
        count,
        fov_threshold=float(cfg_get(cfg, "fov_overlap_threshold", 0.6)),
        min_total_coverage=float(cfg_get(cfg, "min_total_selected_coverage", 0.1)),
        use_plucker=True,
        fov_pool=fov_pool,
    )


def select_memory_indices(
    poses,
    target_positions,
    cfg=None,
    split: str = "training",
    rng=None,
    min_candidate_frame: int = 0,
    latents=None,
    actions=None,
    anchor_candidate_start=None,
    anchor_candidate_stop=None,
    dynamic_stream=None,
):
    """Select memory indices and masks for [anchor][dynamic][revisit]."""

    poses = _as_pose_array(poses)
    target_positions = np.asarray(target_positions, dtype=np.int64)
    target_positions = target_positions[(target_positions >= 0) & (target_positions < len(poses))]
    enabled = bool(cfg_get(cfg, "enabled", True))
    min_candidate_frame = max(0, int(min_candidate_frame))

    counts = {
        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)) if enabled else 0,
        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)) if enabled else 0,
        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)) if enabled else 0,
    }
    indices, masks = _empty_selection(counts)

    if not enabled or len(target_positions) == 0:
        return indices, masks

    target_start = int(target_positions[0])
    policy = _dynamic_policy(cfg)

    anchor_start = min_candidate_frame if anchor_candidate_start is None else max(0, int(anchor_candidate_start))
    if anchor_candidate_stop is None:
        anchor_stop = target_start
    else:
        anchor_stop = min(target_start, max(anchor_start, int(anchor_candidate_stop)))
    anchor_candidates = np.arange(anchor_start, anchor_stop, dtype=np.int64)
    anchor = _select_anchor(anchor_candidates, counts["anchor"], cfg, poses=poses)

    fov_pool = None
    if policy == "multiview" and _dynamic_multiview_selector(cfg) == "fov_greedy":
        fov_pool = _build_shared_fov_candidate_pool(
            poses,
            target_positions,
            cfg,
            split,
            min_candidate_frame=min_candidate_frame,
            dynamic_count=counts["dynamic"],
            revisit_count=counts["revisit"],
        )

    revisit = _select_revisit(
        poses,
        target_positions,
        cfg,
        counts["revisit"],
        np.empty((0,), dtype=np.int64),
        split,
        rng=rng,
        min_candidate_frame=min_candidate_frame,
        fov_pool=fov_pool,
    )
    if policy == "recent":
        dynamic = _select_dynamic_by_policy(
            target_start,
            counts["dynamic"],
            cfg,
            min_candidate_frame=min_candidate_frame,
            target_positions=target_positions,
        )
    elif policy == "event_triggered":
        dynamic_references = revisit if len(revisit) > 0 else target_positions
        if dynamic_stream is not None:
            candidates = _memory_candidate_frames(
                len(poses),
                target_positions,
                cfg,
                split,
                min_candidate_frame=min_candidate_frame,
            )
            candidates = _exclude_local_context(candidates, target_positions, cfg)
            stream = np.asarray(dynamic_stream, dtype=np.int64)
            if len(candidates) > 0:
                eligible_stream = stream[np.isin(stream, candidates)]
            else:
                eligible_stream = np.empty((0,), dtype=np.int64)
            dynamic = _select_dynamic_from_stream(eligible_stream, dynamic_references, counts["dynamic"])
        else:
            dynamic = _select_dynamic_by_policy(
                target_start,
                counts["dynamic"],
                cfg,
                poses=poses,
                latents=latents,
                actions=actions,
                min_candidate_frame=min_candidate_frame,
                reference_frames=dynamic_references,
            )
    elif policy == "multiview":
        dynamic = _select_dynamic_by_policy(
            target_start,
            counts["dynamic"],
            cfg,
            poses=poses,
            min_candidate_frame=min_candidate_frame,
            reference_frames=revisit,
            excluded=revisit,
            target_positions=target_positions,
            split=split,
            fov_pool=fov_pool,
        )
    else:
        raise AssertionError(f"unhandled dynamic selection policy: {policy!r}")

    training_dropout = float(cfg_get(cfg, "training_dropout", 0.0)) if split == "training" else 0.0
    if training_dropout > 0.0 and _rng_random(rng) < training_dropout:
        revisit_candidates = _memory_candidate_frames(
            len(poses),
            target_positions,
            cfg,
            split,
            min_candidate_frame=min_candidate_frame,
        )
        revisit_candidates = _exclude_revisit_local_context(revisit_candidates, target_positions, cfg)
        dynamic_candidates = _dynamic_random_candidates(
            len(poses),
            target_start,
            target_positions,
            cfg,
            split,
            min_candidate_frame,
            policy,
            excluded=revisit if policy == "multiview" else None,
        )
        anchor = _sample_random_selection(anchor_candidates, counts["anchor"], rng=rng)
        dynamic = _sample_random_selection(dynamic_candidates, counts["dynamic"], rng=rng)
        revisit = _sample_random_selection(revisit_candidates, counts["revisit"], rng=rng)

    _write_segment(indices, masks, "anchor", anchor)
    _write_segment(indices, masks, "dynamic", dynamic)
    _write_segment(indices, masks, "revisit", revisit)
    return indices, masks


def memory_segment_lengths(target_length: int, cfg=None) -> Dict[str, int]:
    enabled = bool(cfg_get(cfg, "enabled", True))
    return {
        "target": int(target_length),
        "anchor": int(cfg_get(cfg, "max_anchor_frames", 0)) if enabled else 0,
        "dynamic": int(cfg_get(cfg, "max_dynamic_frames", 0)) if enabled else 0,
        "revisit": int(cfg_get(cfg, "max_revisit_frames", 0)) if enabled else 0,
    }