src/fpgm/config_datagen.py

"""Declarative profile for the DROID -> VACE datagen batch.

Split from :mod:`fpgm.config` deliberately. ``config.py`` holds *algorithm*
tunables (thresholds, resolutions, gate values) that a reviewer changes to alter
what the pipeline computes. This module holds *deployment* configuration -- where
the data lives, which GPUs may be used, how objects are named, where results are
published -- which a reviewer changes to alter where the pipeline runs. Mixing
the two is how ``_DEMO_MESH_PATHS`` ended up as a Python literal inside a stage.

The whole profile loads from one YAML file (``configs/datagen_droid.yaml``) via
:meth:`DatagenProfile.from_yaml`, reusing ``config.py``'s recursive builder so
unknown keys are an error rather than being silently ignored.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import yaml

from fpgm.config import (
    REPO_ROOT,
    DatagenConfig,
    SegmentationConfig,
    TrackingConfig,
    _build,
    _unbuild,
)
from fpgm.types import ConfigError

__all__ = [
    "BatchConfig",
    "DatagenPaths",
    "DatagenProfile",
    "GpuPoolConfig",
    "PublishConfig",
    "RoleConfig",
    "RolesConfig",
    "SceneAssetsConfig",
    "StaticSpanConfig",
]


@dataclass
class DatagenPaths:
    """Every filesystem location the batch reads or writes.

    Relative paths are resolved against the repository root by
    :meth:`DatagenProfile.__post_init__`, so a YAML profile stays portable.
    """

    droid_raw_root: Path = REPO_ROOT / "data" / "droid_raw"
    flows_root: Path = REPO_ROOT / "data" / "pointworld" / "droid" / "flows-fs-optimized"
    cameras_root: Path = REPO_ROOT / "data" / "pointworld" / "droid" / "cameras"
    output_root: Path = REPO_ROOT / "outputs" / "datagen"
    checkpoints_root: Path = REPO_ROOT / "checkpoints"
    urdf: Path = REPO_ROOT / "assets" / "robot" / "pointworld" / "fr3_franka_hand.urdf"
    #: Interpreter used for the one stage that is still a separate process
    #: (S8 export). Declared rather than assumed so a different conda env does
    #: not require a code edit.
    python_executable: Path = Path.home() / "miniconda3" / "envs" / "fpgm" / "bin" / "python"

    def flows_h5(self, uuid: str) -> Path:
        return self.flows_root / f"{uuid}_flows.h5"

    def cameras_json(self, uuid: str) -> Path:
        return self.cameras_root / f"{uuid}_cameras.json"

    def episode_dir(self, uuid: str) -> Path:
        return self.droid_raw_root / uuid

    def camera_dir(self, uuid: str, camera_serial: str) -> Path:
        return self.output_root / uuid / camera_serial

    def master_dir(self, uuid: str, camera_serial: str) -> Path:
        return self.camera_dir(uuid, camera_serial) / "master"


@dataclass
class RoleConfig:
    """How one :class:`~fpgm.datagen.episode_spec.ObjectRole` is treated.

    Keyed by role rather than by object label so it transfers to a scene whose
    nouns are "mug" and "bin" without an edit.
    """

    #: ``"mesh"`` or ``"observed_surface"`` -- see ``ShapeRepresentation``.
    representation: str = "mesh"
    #: ``[lo, hi]`` bound on mesh-to-observation similarity scale, or ``None``.
    plausible_scale_range: list[float] | None = None
    #: Whether pose estimation may re-derive scale from the object mask.
    allow_scale: bool = True
    #: Role whose pose may bridge this role's occlusion gaps, or ``None``.
    #: A manipulated object put *inside* a fixture leaves every camera-facing
    #: mask; the fixture's own pose is the only remaining evidence.
    parent_role: str | None = None


@dataclass
class RolesConfig:
    """Instruction-parsing rules plus the per-role treatment table."""

    #: A derived noun phrase longer than this is not an object name. Measured:
    #: usable DROID phrases here are 1-2 words, unusable ones 6-8.
    max_phrase_words: int = 4
    #: Words that mark a phrase as describing the episode rather than an object.
    meta_tokens: list[str] = field(
        default_factory=lambda: [
            "anything", "task", "tasks", "step", "steps", "consecutively", "suggested",
        ]
    )
    #: ``{uuid: instruction}`` replacing an episode's own ``current_task``.
    #:
    #: **Empty by default, and it should stay that way unless a human has
    #: actually watched the episode.** 4 of the 62 local episodes carry
    #: free-form prompts ("Do anything you like that takes multiple steps to
    #: complete.") that derive no object phrase, so they are refused. They are
    #: recoverable here -- but only as an explicit human assertion about what
    #: the operator did, recorded per uuid and visible in the diff. Guessing
    #: that they *probably* did the same brick/drawer task, and encoding that
    #: guess as a default, is precisely the failure mode this whole module
    #: exists to remove: it would turn 4 unknown episodes into 4
    #: confidently-mislabelled training samples.
    task_overrides: dict[str, str] = field(default_factory=dict)
    manipulated: RoleConfig = field(
        default_factory=lambda: RoleConfig(
            representation="mesh", plausible_scale_range=[0.03, 0.5], allow_scale=True
        )
    )
    fixture: RoleConfig = field(
        default_factory=lambda: RoleConfig(
            representation="observed_surface",
            plausible_scale_range=[0.5, 2.0],
            allow_scale=False,
            parent_role=None,
        )
    )

    def for_role(self, role: Any) -> RoleConfig:
        """``role`` is an ``ObjectRole``; looked up by its ``.value``."""
        name = getattr(role, "value", role)
        cfg = getattr(self, str(name), None)
        if cfg is None:
            raise ConfigError(f"no RoleConfig declared for role {name!r}")
        return cfg


@dataclass
class SceneAssetsConfig:
    """Meshes registered **per scene**, not per episode.

    All 62 local episodes share ``scene_id`` 8756300955, so the manipulated
    object's SAM3D mesh is reconstructed once and reused by every episode in the
    scene. That is a measured property of this dataset, not an assumption: the
    scene id is read from each episode's own flows-h5 root attributes and the
    lookup fails loudly if an episode belongs to a scene with no registered
    mesh.
    """

    #: ``{scene_id: {role_name: mesh_path}}``.
    meshes: dict[str, dict[str, str]] = field(default_factory=dict)

    def mesh_for(self, scene_id: str, role: Any) -> Path:
        name = getattr(role, "value", role)
        by_role = self.meshes.get(str(scene_id))
        if not by_role or str(name) not in by_role:
            raise ConfigError(
                f"no {name!r} mesh registered for scene {scene_id!r}. Either add one under "
                "scene_assets.meshes in the profile, or set that role's representation to "
                "'observed_surface' so its shape is measured from depth instead."
            )
        path = Path(by_role[str(name)])
        if not path.is_absolute():
            path = REPO_ROOT / path
        if not path.exists():
            raise ConfigError(f"registered {name!r} mesh for scene {scene_id!r} not found: {path}")
        return path


@dataclass
class StaticSpanConfig:
    """Thresholds for measuring an object's leading static run.

    Replaces the hard-coded ``(0, 8)``. See
    :class:`fpgm.datagen.static_span.StaticSpanEstimator` for the measurement.
    """

    #: Per-frame centroid displacement (pixels) below which a frame counts as
    #: static. Scaled to the native video width by the estimator.
    max_step_px: float = 1.5
    #: Total drift (pixels) the whole span may accumulate. A slow, steady creep
    #: passes the per-step test frame by frame while being real motion.
    max_total_drift_px: float = 4.0
    #: Shortest run that is worth reporting. Below this the pose-noise estimate
    #: it feeds is dominated by sample count, not by noise.
    min_frames: int = 4
    #: Longest span to return. A very long "static" run usually means the mask
    #: is tracking something that never moves at all, and averaging over it
    #: hides real early motion.
    max_frames: int = 30


@dataclass
class GpuPoolConfig:
    """Which GPUs the batch may use, and how densely.

    The box has 4xH200 shared with other users whose free VRAM was observed to
    swing by tens of GB within 90 s. Admission is therefore by *measured free
    memory at claim time*, not by a static device list.
    """

    #: Explicit device ids to consider, or ``None`` to consider all visible.
    device_ids: list[int] | None = None
    #: A device with less free VRAM than this is not offered to a worker.
    min_free_mb: int = 16000
    #: Concurrent episodes per device. One SAM3 + one TAPNext is ~10-14 GB, so
    #: 1 is the safe default on a contended box; raise it only after measuring.
    workers_per_device: int = 1
    #: Hard cap on concurrent workers across all devices, or ``None`` for
    #: ``len(devices) * workers_per_device``. Bounds CPU/IO pressure
    #: independently of GPU count.
    max_workers: int | None = None
    #: Threads each worker's numeric libraries may use. Unset, 128 cores x N
    #: workers of BLAS threads thrash; this is the "don't spam the CPU" knob.
    threads_per_worker: int = 8
    #: How long a worker waits for a device to free up before failing.
    acquire_timeout_s: float = 3600.0
    poll_interval_s: float = 20.0


@dataclass
class BatchConfig:
    """Batch-level policy: what to run, in what order, how to report."""

    camera_role: str = "ext1"
    #: Fraction of episodes in the pilot pass. The pilot runs first, is
    #: inspected, and only then does the remainder run.
    pilot_fraction: float = 0.10
    #: Minimum pilot episodes, so a small dataset still gets a real pilot.
    pilot_min_episodes: int = 4
    #: Fraction of pilot episodes that must reach ``ok`` before the full batch
    #: is allowed to start automatically.
    pilot_pass_threshold: float = 0.6
    #: Per-episode wall-clock budget; a hung episode must not stall the batch.
    episode_timeout_s: float = 5400.0
    #: Vacated-region hole fraction above which a window is flagged for human
    #: attention as a genuine single-viewpoint capture limitation.
    high_hole_fraction_gate: float = 0.10
    #: Delete the large per-episode intermediates (``robot_buffers/*.npy``,
    #: ~336 MB/episode) once the episode's exports exist. 62 episodes would
    #: otherwise cost ~21 GB on a filesystem already at 99%.
    prune_intermediates: bool = True
    #: How S4 gets its object identities: ``"text"`` (default) resolves
    #: :class:`~fpgm.datagen.episode_spec.EpisodeSpecResolver`'s instruction
    #: parse first and only falls back to motion clustering per-role if a
    #: role's every SAM3 text prompt matches nothing (unchanged from before
    #: this setting existed -- see
    #: :meth:`~fpgm.datagen.pipeline.EpisodePipeline._run_s4`'s own
    #: docstring). ``"motion"`` runs
    #: :class:`~fpgm.datagen.discovery.DynamicObjectStage` FIRST, unconditionally,
    #: and never attempts a text prompt at all -- object identity for the whole
    #: episode is motion-only, and an unparseable DROID instruction is no
    #: longer a reason to skip the episode. See
    #: :mod:`fpgm.datagen.episode_spec`'s module docstring for the full
    #: rationale.
    objects_from: str = "text"

    def __post_init__(self) -> None:
        if self.objects_from not in ("text", "motion"):
            raise ConfigError(
                f"batch.objects_from must be 'text' or 'motion', got {self.objects_from!r}"
            )


@dataclass
class PublishConfig:
    """Where the finished dataset is uploaded, and what goes in it."""

    #: Destination bucket URL. Buckets are **not** repos: ``repo_info`` rejects
    #: ``repo_type="bucket"`` (valid types are only model/dataset/space) and
    #: returns 404 for this id as a dataset. The Python API for them is
    #: ``HfApi.sync_bucket(source=<local dir>, dest=<this url>)``, which has its
    #: own ``dry_run`` and ``plan`` modes -- verified live against
    #: ``huggingface_hub`` 1.25.1 on this box.
    bucket_url: str = "hf://buckets/twanghcmut/droid-processing"
    #: Per-episode ``master/`` files worth shipping. The 336 MB of
    #: ``robot_buffers/*.npy`` per episode are deliberately absent: they are
    #: reproducible from the URDF plus the trajectory in seconds, and shipping
    #: them would make the dataset ~21 GB instead of ~1.5 GB.
    master_includes: list[str] = field(
        default_factory=lambda: [
            "poses.npz",
            "events.json",
            "episode_spec.json",
            "background_depth.h5",
            "background_depth.png",
            "plate.png",
            "meshes/**",
            "*_masks.h5",
            "robot_buffers/meta.json",
            "robot_buffers/robot_seg.mkv",
        ]
    )
    #: Per-window ``vace/`` files. This is the training payload.
    vace_includes: list[str] = field(
        default_factory=lambda: [
            "target.mp4",
            "control_depth.mkv",
            "control_seg.mkv",
            "control_normal.mp4",
            "fg_mask.mp4",
            "ref_plate.png",
            "caption.txt",
            "sample.json",
        ]
    )
    #: Also ship the eyeball-checkable debug renders. Small, and the project's
    #: own lesson was that debug renders outrank metric tables.
    include_debug_videos: bool = True


@dataclass
class DatagenProfile:
    """The whole deployment profile. One YAML file, one object."""

    paths: DatagenPaths = field(default_factory=DatagenPaths)
    roles: RolesConfig = field(default_factory=RolesConfig)
    scene_assets: SceneAssetsConfig = field(default_factory=SceneAssetsConfig)
    static_span: StaticSpanConfig = field(default_factory=StaticSpanConfig)
    gpu: GpuPoolConfig = field(default_factory=GpuPoolConfig)
    batch: BatchConfig = field(default_factory=BatchConfig)
    publish: PublishConfig = field(default_factory=PublishConfig)
    #: Algorithm thresholds, unchanged from ``config.py``'s own defaults unless
    #: the profile overrides them.
    datagen: DatagenConfig = field(default_factory=DatagenConfig)
    #: SAM 3.1 settings. Lives in the profile rather than being constructed
    #: inline in a stage because it is the pipeline's dominant VRAM consumer
    #: and this box is shared -- the memory/speed trade has to be tunable
    #: without a code edit. Also gives every episode in a shard the *same*
    #: config object, which is what lets ``ModelRegistry`` reuse one predictor
    #: across the whole shard.
    segmentation: SegmentationConfig = field(default_factory=SegmentationConfig)
    #: TAPNext++ settings, for the same reason.
    tracking: TrackingConfig = field(default_factory=TrackingConfig)

    def __post_init__(self) -> None:
        for name in ("droid_raw_root", "flows_root", "cameras_root", "output_root",
                     "checkpoints_root", "urdf", "python_executable"):
            value = getattr(self.paths, name)
            if not Path(value).is_absolute():
                setattr(self.paths, name, REPO_ROOT / value)

    @classmethod
    def from_yaml(cls, path: str | Path) -> DatagenProfile:
        with open(path) as fh:
            raw = yaml.safe_load(fh) or {}
        return _build(cls, raw)

    def to_yaml(self, path: str | Path) -> None:
        with open(path, "w") as fh:
            yaml.safe_dump(_unbuild(self), fh, sort_keys=False)