src/data_contract.py

"""Backbone-neutral trajectory contract for Route A.

The module deliberately stops before model-specific normalization. Every backbone
therefore receives the same selected RGB frames, labels and masks.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any, Iterable, Literal, Sequence

import numpy as np
import torch
from PIL import Image


TemporalStrategy = Literal["forward", "reverse", "rewind", "different_task"]


def deterministic_frame_indices(
    frame_count: int, target_frames: int = 8
) -> tuple[np.ndarray, np.ndarray]:
    """Return official-style rint(linspace) indices and a real-frame mask."""

    if frame_count <= 0 or target_frames <= 0:
        raise ValueError("frame_count and target_frames must be positive")
    if frame_count >= target_frames:
        indices = np.rint(
            np.linspace(0, frame_count - 1, target_frames)
        ).astype(np.int64)
        indices[0] = 0
        indices[-1] = frame_count - 1
        return indices, np.ones(target_frames, dtype=np.bool_)

    real = np.arange(frame_count, dtype=np.int64)
    padding = np.full(target_frames - frame_count, frame_count - 1, dtype=np.int64)
    valid = np.concatenate(
        [
            np.ones(frame_count, dtype=np.bool_),
            np.zeros(target_frames - frame_count, dtype=np.bool_),
        ]
    )
    return np.concatenate([real, padding]), valid


def temporal_permutation(length: int, strategy: TemporalStrategy) -> np.ndarray:
    """Return a permutation shared by frames, labels, indices and masks."""

    if length <= 0:
        raise ValueError("length must be positive")
    if strategy in ("forward", "different_task"):
        return np.arange(length, dtype=np.int64)
    if strategy == "reverse":
        return np.arange(length - 1, -1, -1, dtype=np.int64)
    if strategy == "rewind":
        split = length // 2
        return np.concatenate(
            [
                np.arange(split, dtype=np.int64),
                np.arange(length - 1, split - 1, -1, dtype=np.int64),
            ]
        )
    raise ValueError(f"unsupported temporal strategy: {strategy}")


def letterbox_rgb(frame: np.ndarray, canvas_size: int = 256) -> np.ndarray:
    """Deterministically resize an RGB uint8 frame onto a square black canvas."""

    if (
        frame.ndim != 3
        or frame.shape[-1] != 3
        or frame.dtype != np.uint8
        or canvas_size <= 0
    ):
        raise ValueError("frame must be uint8 [H,W,3] and canvas_size must be positive")
    height, width = frame.shape[:2]
    if height <= 0 or width <= 0:
        raise ValueError("frame dimensions must be positive")
    scale = min(canvas_size / width, canvas_size / height)
    resized_width = max(1, min(canvas_size, round(width * scale)))
    resized_height = max(1, min(canvas_size, round(height * scale)))
    resized = Image.fromarray(frame, mode="RGB").resize(
        (resized_width, resized_height), Image.Resampling.BILINEAR
    )
    canvas = Image.new("RGB", (canvas_size, canvas_size), (0, 0, 0))
    offset = (
        (canvas_size - resized_width) // 2,
        (canvas_size - resized_height) // 2,
    )
    canvas.paste(resized, offset)
    return np.asarray(canvas, dtype=np.uint8)


def _c51_targets(values: np.ndarray, bins: int) -> np.ndarray:
    if bins < 2:
        raise ValueError("progress bins must be at least 2")
    positions = np.clip(values, 0.0, 1.0) * (bins - 1)
    lower = np.floor(positions).astype(np.int64)
    upper = np.ceil(positions).astype(np.int64)
    output = np.zeros((len(values), bins), dtype=np.float32)
    rows = np.arange(len(values))
    output[rows, lower] += upper - positions
    output[rows, upper] += positions - lower
    exact = lower == upper
    output[rows[exact], lower[exact]] = 1.0
    return output


@dataclass(frozen=True)
class TrajectoryContract:
    trajectory_id: str
    task: str
    data_source: str
    quality_label: str
    frames_uint8: np.ndarray
    source_frame_indices: np.ndarray
    frame_valid_mask: np.ndarray
    progress: np.ndarray
    progress_distribution: np.ndarray
    progress_loss_mask: np.ndarray
    success_labels: np.ndarray
    success_loss_mask: np.ndarray
    strategy: TemporalStrategy
    metadata: dict[str, Any]

    def validate(self) -> None:
        frame_count = self.frames_uint8.shape[0]
        if self.frames_uint8.dtype != np.uint8 or self.frames_uint8.ndim != 4:
            raise ValueError("frames_uint8 must be uint8 [T,H,W,3]")
        if self.frames_uint8.shape[-1] != 3:
            raise ValueError("frames_uint8 must have three RGB channels")
        aligned = (
            self.source_frame_indices,
            self.frame_valid_mask,
            self.progress,
            self.progress_distribution,
            self.progress_loss_mask,
            self.success_labels,
            self.success_loss_mask,
        )
        if any(value.shape[0] != frame_count for value in aligned):
            raise ValueError("all frame-level fields must have the same leading size")
        if np.any(self.progress < 0) or np.any(self.progress > 1):
            raise ValueError("progress values must be in [0,1]")
        if np.any(self.progress_loss_mask & ~self.frame_valid_mask):
            raise ValueError("padding may not enter progress loss")
        if np.any(self.success_loss_mask & ~self.frame_valid_mask):
            raise ValueError("padding may not enter success loss")


@dataclass(frozen=True)
class PreferenceContract:
    trajectory_a: TrajectoryContract
    trajectory_b: TrajectoryContract
    preference_label: float
    strategy: str

    def swapped(self) -> PreferenceContract:
        return replace(
            self,
            trajectory_a=self.trajectory_b,
            trajectory_b=self.trajectory_a,
            preference_label=1.0 - self.preference_label,
        )


def prepare_trajectory(
    *,
    trajectory_id: str,
    frames: np.ndarray,
    task: str,
    data_source: str,
    quality_label: str,
    partial_success: float | None = None,
    strategy: TemporalStrategy = "forward",
    target_frames: int = 8,
    canvas_size: int = 256,
    progress_bins: int = 10,
    success_cutoff: float = 1.0,
    metadata: dict[str, Any] | None = None,
) -> TrajectoryContract:
    """Materialize one auditable trajectory from an official sampler trajectory."""

    if frames.ndim != 4 or frames.shape[-1] != 3 or frames.dtype != np.uint8:
        raise ValueError("frames must be uint8 [T,H,W,3]")
    if not trajectory_id or not task or not data_source:
        raise ValueError("trajectory_id, task and data_source are required")
    if partial_success is not None and not 0.0 <= partial_success <= 1.0:
        raise ValueError("partial_success must be in [0,1]")
    if not 0.0 <= success_cutoff <= 1.0:
        raise ValueError("success_cutoff must be in [0,1]")

    selected_indices, frame_valid = deterministic_frame_indices(
        len(frames), target_frames
    )
    permutation = temporal_permutation(target_frames, strategy)
    selected_indices = selected_indices[permutation]
    frame_valid = frame_valid[permutation]
    selected_frames = np.stack(
        [letterbox_rgb(frames[index], canvas_size) for index in selected_indices]
    )

    denominator = max(len(frames) - 1, 1)
    progress = selected_indices.astype(np.float32) / denominator
    if partial_success is not None:
        progress *= partial_success
    if strategy == "different_task":
        progress.fill(0.0)

    quality = quality_label.lower()
    dense_progress_allowed = quality not in {"failure", "failed", "suboptimal"}
    progress_loss_mask = frame_valid & dense_progress_allowed
    if strategy == "different_task":
        progress_loss_mask = frame_valid.copy()
    success_labels = (
        (progress >= success_cutoff)
        & (quality == "successful")
        & frame_valid
    ).astype(np.float32)
    success_loss_mask = frame_valid.copy()
    audit_metadata = dict(metadata or {})
    audit_metadata.update(
        {
            "source_frame_count": int(len(frames)),
            "selected_frame_indices": selected_indices.tolist(),
            "frame_valid_mask": frame_valid.tolist(),
            "temporal_strategy": strategy,
            "canvas_size": canvas_size,
            "letterbox_fill_rgb": [0, 0, 0],
        }
    )
    result = TrajectoryContract(
        trajectory_id=trajectory_id,
        task=task,
        data_source=data_source,
        quality_label=quality_label,
        frames_uint8=selected_frames,
        source_frame_indices=selected_indices,
        frame_valid_mask=frame_valid,
        progress=progress,
        progress_distribution=_c51_targets(progress, progress_bins),
        progress_loss_mask=progress_loss_mask,
        success_labels=success_labels,
        success_loss_mask=success_loss_mask,
        strategy=strategy,
        metadata=audit_metadata,
    )
    result.validate()
    return result


def adapt_official_trajectory(
    trajectory: Any,
    *,
    strategy: TemporalStrategy,
    progress_supervised: bool,
    canvas_size: int = 256,
    progress_bins: int = 10,
) -> TrajectoryContract:
    """Adapt an official Robometer sampler trajectory without changing its labels.

    The sampler must expose its selected source indices and padding mask in
    ``trajectory.metadata``. This is intentionally strict: silently re-sampling
    here would invalidate the matched-backbone comparison.
    """

    metadata = dict(trajectory.metadata or {})
    if "selected_frame_indices" not in metadata or "frame_valid_mask" not in metadata:
        raise ValueError("official sampler metadata lacks frozen frame indices/mask")
    frames = np.asarray(trajectory.frames)
    if frames.ndim != 4 or frames.shape[-1] != 3 or frames.dtype != np.uint8:
        raise ValueError("official sampler frames must be uint8 [T,H,W,3]")
    selected_indices = np.asarray(metadata["selected_frame_indices"], dtype=np.int64)
    frame_valid = np.asarray(metadata["frame_valid_mask"], dtype=np.bool_)
    if len(selected_indices) != len(frames) or len(frame_valid) != len(frames):
        raise ValueError("official sampler frame metadata is not aligned")

    target_progress = trajectory.target_progress
    if target_progress is None or len(target_progress) != len(frames):
        raise ValueError("official target_progress is missing or misaligned")
    first_target = target_progress[0]
    if isinstance(first_target, torch.Tensor) or np.asarray(first_target).ndim > 0:
        progress_distribution = np.stack(
            [np.asarray(value, dtype=np.float32) for value in target_progress]
        )
        if progress_distribution.shape != (len(frames), progress_bins):
            raise ValueError("official discrete progress target has unexpected shape")
        support = np.linspace(0.0, 1.0, progress_bins, dtype=np.float32)
        progress = progress_distribution @ support
    else:
        progress = np.asarray(target_progress, dtype=np.float32)
        progress_distribution = _c51_targets(progress, progress_bins)

    predict_mask = trajectory.predict_last_frame_mask
    if predict_mask is None:
        predict_valid = np.ones(len(frames), dtype=np.bool_)
    else:
        predict_valid = np.asarray(predict_mask, dtype=np.float32) > 0
    progress_loss_mask = frame_valid & predict_valid & progress_supervised
    success_labels = np.asarray(trajectory.success_label, dtype=np.float32)
    if success_labels.shape != (len(frames),):
        raise ValueError("official success labels are missing or misaligned")

    metadata.update(
        {
            "adapter": "official_robometer_sampler_v1",
            "temporal_strategy": strategy,
            "canvas_size": canvas_size,
            "letterbox_fill_rgb": [0, 0, 0],
        }
    )
    result = TrajectoryContract(
        trajectory_id=trajectory.id,
        task=trajectory.task,
        data_source=trajectory.data_source,
        quality_label=trajectory.quality_label,
        frames_uint8=np.stack(
            [letterbox_rgb(frame, canvas_size) for frame in frames]
        ),
        source_frame_indices=selected_indices,
        frame_valid_mask=frame_valid,
        progress=progress,
        progress_distribution=progress_distribution,
        progress_loss_mask=progress_loss_mask,
        success_labels=success_labels,
        success_loss_mask=frame_valid.copy(),
        strategy=strategy,
        metadata=metadata,
    )
    result.validate()
    return result


def make_preference_contract(
    chosen: TrajectoryContract,
    rejected: TrajectoryContract,
    *,
    strategy: str,
    chosen_first: bool = True,
) -> PreferenceContract:
    if chosen_first:
        return PreferenceContract(chosen, rejected, 1.0, strategy)
    return PreferenceContract(rejected, chosen, 0.0, strategy)


class CosmosRBMCollator:
    """Stack the frozen contract without invoking a backbone processor."""

    schema_version = "cosmos_rbm_batch_v1"

    @staticmethod
    def _stack_trajectories(items: Sequence[TrajectoryContract]) -> dict[str, Any]:
        for item in items:
            item.validate()
        return {
            "frames_uint8": torch.from_numpy(
                np.stack([item.frames_uint8 for item in items])
            ),
            "source_frame_indices": torch.from_numpy(
                np.stack([item.source_frame_indices for item in items])
            ),
            "frame_valid_mask": torch.from_numpy(
                np.stack([item.frame_valid_mask for item in items])
            ),
            "progress": torch.from_numpy(np.stack([item.progress for item in items])),
            "progress_distribution": torch.from_numpy(
                np.stack([item.progress_distribution for item in items])
            ),
            "progress_loss_mask": torch.from_numpy(
                np.stack([item.progress_loss_mask for item in items])
            ),
            "success_labels": torch.from_numpy(
                np.stack([item.success_labels for item in items])
            ),
            "success_loss_mask": torch.from_numpy(
                np.stack([item.success_loss_mask for item in items])
            ),
            "trajectory_ids": [item.trajectory_id for item in items],
            "tasks": [item.task for item in items],
            "data_sources": [item.data_source for item in items],
            "quality_labels": [item.quality_label for item in items],
            "strategies": [item.strategy for item in items],
            "metadata": [item.metadata for item in items],
        }

    def __call__(
        self, samples: Sequence[TrajectoryContract | PreferenceContract]
    ) -> dict[str, Any]:
        trajectories = [item for item in samples if isinstance(item, TrajectoryContract)]
        preferences = [item for item in samples if isinstance(item, PreferenceContract)]
        output: dict[str, Any] = {
            "schema_version": self.schema_version,
            "num_trajectories": len(trajectories),
            "num_preferences": len(preferences),
        }
        if trajectories:
            output["trajectory_inputs"] = self._stack_trajectories(trajectories)
        if preferences:
            output["preference_inputs"] = {
                "trajectory_a": self._stack_trajectories(
                    [item.trajectory_a for item in preferences]
                ),
                "trajectory_b": self._stack_trajectories(
                    [item.trajectory_b for item in preferences]
                ),
                "preference_labels": torch.tensor(
                    [item.preference_label for item in preferences], dtype=torch.float32
                ),
                "strategies": [item.strategy for item in preferences],
            }
        return output


def validate_span_bounds(
    bounds: Sequence[tuple[int, int]], frame_count: int, sequence_length: int
) -> None:
    """Validate the one-frame/one-nonoverlapping-token-span contract."""

    if len(bounds) != frame_count:
        raise ValueError("visual span count must match frame count")
    previous_stop = -1
    for start, stop in bounds:
        if start < 0 or stop <= start or stop > sequence_length:
            raise ValueError("visual token span is empty or out of bounds")
        if start < previous_stop:
            raise ValueError("visual token spans overlap")
        previous_stop = stop


def worker_invariant_sample_ids(
    ids: Iterable[str], *, seed: int, worker_count: int
) -> list[str]:
    """Simulate round-robin worker sharding while preserving global sample order."""

    if worker_count <= 0:
        raise ValueError("worker_count must be positive")
    materialized = list(ids)
    unique = list(dict.fromkeys(materialized))
    if len(unique) != len(materialized):
        raise ValueError("sample ids must be unique")
    ranked = sorted(
        unique,
        key=lambda value: hashlib.sha256(f"{seed}:{value}".encode()).digest(),
    )
    shards = [list(enumerate(ranked))[worker::worker_count] for worker in range(worker_count)]
    return [value for _, value in sorted(item for shard in shards for item in shard)]


def audit_group_overlap(train_ids: Iterable[str], eval_ids: Iterable[str]) -> set[str]:
    return set(train_ids) & set(eval_ids)


def dump_batch_audit(batch: dict[str, Any], output: Path) -> None:
    """Write metadata and tensor shapes only; never dump trajectory pixels."""

    def describe(value: Any) -> Any:
        if isinstance(value, torch.Tensor):
            return {"shape": list(value.shape), "dtype": str(value.dtype)}
        if isinstance(value, dict):
            return {key: describe(item) for key, item in value.items()}
        return value

    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(describe(batch), indent=2) + "\n", encoding="utf-8")