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")
_FOV_NUM_POINTS = 10000
_FOV_RADIUS = 30.0
_FOV_HALF_H = 105.0 / 2.0
_FOV_HALF_V = 75.0 / 2.0
_POSE_DISTANCE_SCALE = 30.0
_ANGLE_DISTANCE_SCALE = 180.0


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_revisit_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 _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 _sample_points_in_sphere(center: torch.Tensor) -> torch.Tensor:
    device = center.device
    dtype = center.dtype

    samples_r = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
    samples_phi = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
    samples_u = torch.rand(_FOV_NUM_POINTS, device=device, dtype=dtype)
    r = _FOV_RADIUS * torch.pow(samples_r, 1.0 / 3.0)
    phi = 2.0 * torch.pi * samples_phi
    theta = torch.acos(1.0 - 2.0 * samples_u)

    offsets = torch.stack(
        [
            r * torch.sin(theta) * torch.cos(phi),
            r * torch.sin(theta) * torch.sin(phi),
            r * torch.cos(theta),
        ],
        dim=-1,
    )
    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]
    azimuth = torch.rad2deg(torch.atan2(x, z))
    elevation = torch.rad2deg(torch.atan2(y, torch.sqrt(x * x + z * z)))
    yaw_diff = _wrap_degrees(azimuth - poses[:, None, 4])
    pitch_diff = _wrap_degrees(elevation - poses[:, None, 3])
    return (yaw_diff < _FOV_HALF_H) & (pitch_diff < _FOV_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) -> np.ndarray:
    topk = cfg_get(cfg, "pose_preselect_topk", 64)
    if topk is None:
        return candidates
    topk = int(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", 64))
    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_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 _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 _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)

    poses_t = _as_pose_tensor(poses)
    candidates_t = torch.as_tensor(candidates.astype(np.int64), device=poses_t.device, dtype=torch.long)
    targets_t = torch.as_tensor(target_positions.astype(np.int64), 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)

    # Training preselects with cheap pose similarity; Plucker ranking keeps it
    # closer to the inference FOV selector without running the full 10k-point path.
    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()):
        return np.empty((0,), dtype=np.int64)

    valid_idx = torch.nonzero(valid, as_tuple=False).flatten()
    valid_candidates = candidates_t.index_select(0, valid_idx)
    valid_scores = pose_scores.index_select(0, valid_idx)

    if bool(cfg_get(cfg, "training_use_plucker", True)):
        plucker = _plucker_scores_tensor(
            poses_t.index_select(0, valid_candidates),
            target_poses,
            cfg,
        )
        rank_scores = valid_scores + float(cfg_get(cfg, "training_plucker_weight", 1.0)) * plucker
        order = torch.argsort(rank_scores, descending=True)[:count]
        selected = valid_candidates.index_select(0, order).cpu().numpy()
    else:
        valid_np = valid_candidates.cpu().numpy()
        size = min(count, len(valid_np))
        choice = np.random.choice if rng is None else rng.choice
        selected = choice(valid_np, size=size, replace=False)

    return np.sort(np.asarray(selected, dtype=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,
) -> np.ndarray:
    if count <= 0 or len(candidates) == 0 or len(target_positions) == 0:
        return np.empty((0,), dtype=np.int64)

    poses_t = _as_pose_tensor(poses)
    candidates = _pose_preselect(candidates.astype(np.int64), poses_t, target_positions, cfg)
    inside, fov_values = _candidate_fov_masks(poses_t, candidates, target_positions, cfg)
    if inside.shape[0] == 0:
        return np.empty((0,), dtype=np.int64)

    valid = torch.ones((len(candidates),), device=poses_t.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 = torch.as_tensor(candidates, device=poses_t.device, dtype=torch.long).index_select(0, valid_idx)
    inside = inside.index_select(0, valid_idx)
    fov_values = fov_values.index_select(0, valid_idx)
    target_t = torch.as_tensor(target_positions, device=poses_t.device, dtype=torch.long)
    if use_plucker:
        plucker = _plucker_scores_tensor(poses_t.index_select(0, candidates_t), poses_t.index_select(0, target_t), cfg)
    else:
        plucker = torch.zeros((candidates_t.shape[0],), device=poses_t.device, dtype=torch.float32)

    remaining = torch.ones((candidates_t.shape[0],), device=poses_t.device, dtype=torch.bool)
    covered = torch.zeros((inside.shape[1],), device=poses_t.device, dtype=torch.bool)
    selected_rows = []
    first_target = int(target_positions[0])
    gaps = first_target - candidates_t

    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=poses_t.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", {})
    return str(cfg_get(dynamic_cfg, "selection_policy", "recent"))


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,
    )
    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,
) -> np.ndarray:
    policy = _dynamic_policy(cfg)
    if policy == "recent":
        return _select_dynamic(target_start, 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,
        )
    raise ValueError(f"memory_selection.dynamic.selection_policy must be recent or event_triggered; got {policy!r}")


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,
) -> 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,
    )


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)
    if policy not in {"recent", "event_triggered"}:
        raise ValueError(f"memory_selection.dynamic.selection_policy must be recent or event_triggered; got {policy!r}")

    if policy == "recent":
        dynamic = _select_dynamic_by_policy(
            target_start,
            counts["dynamic"],
            cfg,
            min_candidate_frame=min_candidate_frame,
        )
    else:
        dynamic = np.empty((0,), dtype=np.int64)

    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)

    revisit = _select_revisit(
        poses,
        target_positions,
        cfg,
        counts["revisit"],
        np.empty((0,), dtype=np.int64),
        split,
        rng=rng,
        min_candidate_frame=min_candidate_frame,
    )

    if 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,
            )
            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,
            )

    _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,
    }