src/mcss/training/supervision.py

"""Training-private query supervision with strict train-scene identity checks."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch
from PIL import Image
from torch import Tensor

from mcss.types import Cameras

QUERY_VAULT_SCHEMA = "mcss.dynamic.query_vault.v1"
_EPISODE_FIELDS = frozenset(
    {
        "schema_version",
        "episode_id",
        "scene_id",
        "split_id",
        "query_vault_id",
        "image_size",
        "query",
    }
)
_QUERY_FIELDS = frozenset({"frame_id", "intrinsics", "c2w", "rgb", "depth"})


@dataclass(frozen=True)
class TrainingSupervision:
    """Raw-world query cameras and labels available only to offline training."""

    episode_id: str
    scene_id: str
    query_vault_id: str
    frame_ids: tuple[int, ...]
    query_cameras: Cameras
    query_rgb: Tensor
    query_depth: Tensor

    @classmethod
    def from_query_file(
        cls,
        path: str | Path,
        *,
        episode_id: str,
        scene_id: str,
        query_vault_id: str,
        context_frame_ids: tuple[int, ...],
        device: str | torch.device = "cpu",
    ) -> TrainingSupervision:
        """Open one train-only query artifact after identity and overlap validation."""

        artifact = _read_query_artifact(path)
        _validate_identity(artifact, episode_id, scene_id, query_vault_id)
        image_size = _validate_image_size(artifact["image_size"])
        records = _validate_query_records(artifact["query"], image_size, context_frame_ids)
        target_device = torch.device(device)
        cameras = _query_cameras(records, image_size, target_device)
        rgb = torch.stack(
            [_load_rgb(record["rgb"], image_size, target_device) for record in records]
        )
        depth = torch.stack(
            [_load_depth(record["depth"], image_size, target_device) for record in records]
        )
        return cls(
            episode_id=episode_id,
            scene_id=scene_id,
            query_vault_id=query_vault_id,
            frame_ids=tuple(record["frame_id"] for record in records),
            query_cameras=cameras,
            query_rgb=rgb.unsqueeze(0),
            query_depth=depth.unsqueeze(0),
        )


def _read_query_artifact(path: str | Path) -> dict[str, Any]:
    source = Path(path)
    try:
        artifact = json.loads(source.read_text(encoding="utf-8"))
    except json.JSONDecodeError as error:
        raise ValueError(f"invalid training query JSON: {source}") from error
    if not isinstance(artifact, dict):
        raise ValueError("training query JSON root must be an object")
    _require_exact_fields(artifact, _EPISODE_FIELDS, "training query artifact")
    if artifact["schema_version"] != QUERY_VAULT_SCHEMA:
        raise ValueError(f"expected training query schema {QUERY_VAULT_SCHEMA}")
    return artifact


def _validate_identity(
    artifact: dict[str, Any], episode_id: str, scene_id: str, query_vault_id: str
) -> None:
    for name, expected in {
        "episode_id": episode_id,
        "scene_id": scene_id,
        "query_vault_id": query_vault_id,
    }.items():
        if not isinstance(expected, str) or not expected:
            raise ValueError(f"expected {name} must be a nonempty string")
        if artifact[name] != expected:
            raise ValueError(f"training query {name} does not match the online episode")
    if artifact["split_id"] != "train":
        raise ValueError("training supervision requires a training split query artifact")


def _validate_query_records(
    value: Any, image_size: tuple[int, int], context_frame_ids: tuple[int, ...]
) -> list[dict[str, Any]]:
    if not isinstance(value, list) or not value:
        raise ValueError("training query artifact requires a nonempty query list")
    context_ids = set(context_frame_ids)
    if len(context_ids) != len(context_frame_ids) or any(
        type(frame_id) is not int or frame_id < 0 for frame_id in context_frame_ids
    ):
        raise ValueError("context frame IDs must be unique nonnegative integers")
    records: list[dict[str, Any]] = []
    seen: set[int] = set()
    for record in value:
        if not isinstance(record, dict):
            raise ValueError("training query records must be objects")
        _require_exact_fields(record, _QUERY_FIELDS, "training query frame")
        frame_id = record["frame_id"]
        if type(frame_id) is not int or frame_id < 0:
            raise ValueError("training query frame IDs must be nonnegative integers")
        if frame_id in seen:
            raise ValueError("training query frame IDs must be unique")
        if frame_id in context_ids:
            raise ValueError("training query frame IDs must be disjoint from context")
        seen.add(frame_id)
        for name in ("rgb", "depth"):
            if not isinstance(record[name], str) or not record[name]:
                raise ValueError(f"training query {name} reference must be a nonempty path string")
        _camera_from_record(record, image_size, torch.device("cpu"))
        records.append(record)
    if [record["frame_id"] for record in records] != sorted(seen):
        raise ValueError("training query frame IDs must be in ascending order")
    return records


def _query_cameras(
    records: list[dict[str, Any]], image_size: tuple[int, int], device: torch.device
) -> Cameras:
    cameras = [_camera_from_record(record, image_size, device) for record in records]
    return Cameras(
        torch.stack([camera.intrinsics for camera in cameras]).unsqueeze(0),
        torch.stack([camera.c2w for camera in cameras]).unsqueeze(0),
        image_size,
    )


def _camera_from_record(
    record: dict[str, Any], image_size: tuple[int, int], device: torch.device
) -> Cameras:
    try:
        intrinsics = torch.tensor(record["intrinsics"], dtype=torch.float32, device=device)
        c2w = torch.tensor(record["c2w"], dtype=torch.float32, device=device)
    except (TypeError, ValueError) as error:
        raise ValueError("training query cameras must contain numeric matrices") from error
    return Cameras(intrinsics, c2w, image_size)


def _load_rgb(path: str, image_size: tuple[int, int], device: torch.device) -> Tensor:
    source = Path(path)
    with Image.open(source) as image:
        image = image.convert("RGB")
        if (image.height, image.width) != image_size:
            raise ValueError(f"training query RGB size does not match image_size: {source}")
        array = np.asarray(image, dtype=np.float32).copy()
    return torch.from_numpy(array).permute(2, 0, 1).contiguous().div(255.0).to(device)


def _load_depth(path: str, image_size: tuple[int, int], device: torch.device) -> Tensor:
    value = np.load(Path(path), allow_pickle=False)
    if value.shape != image_size:
        raise ValueError("training query depth size does not match image_size")
    return torch.from_numpy(np.asarray(value, dtype=np.float32).copy()).unsqueeze(0).to(device)


def _validate_image_size(value: Any) -> tuple[int, int]:
    if (
        not isinstance(value, list)
        or len(value) != 2
        or any(type(item) is not int or item < 2 for item in value)
    ):
        raise ValueError(
            "training query image_size must be [height, width] with values at least two"
        )
    return value[0], value[1]


def _require_exact_fields(value: dict[str, Any], expected: frozenset[str], name: str) -> None:
    actual = frozenset(value)
    if actual != expected:
        missing = sorted(expected - actual)
        unexpected = sorted(actual - expected)
        raise ValueError(f"{name} fields mismatch; missing={missing}, unexpected={unexpected}")