ylff/models/intermediate_artifacts.py

"""
Canonical intermediate artifact schemas.

These types are the stable "contract" between:
- ingest (raw bundle -> validated inputs),
- teacher (depth + σ + provenance),
- audit/calibration (measurement-level results + calibration params),
- training (datasets + checkpoints),
- inference (outputs + diagnostics).

Important: large arrays are referenced by URI/path; we do NOT embed dense tensors
in JSON. Use an ArtifactStore to persist and reference big payloads.
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Tuple
from pydantic import BaseModel, Field

from .spec_enums import OperatingRegime


class IntermediateSchemaVersion(str, Enum):
    V1_0 = "1.0"


class Units(str, Enum):
    METERS = "meters"
    PIXELS = "pixels"
    SECONDS = "seconds"
    NONE = "none"


class ArtifactURI(BaseModel):
    """
    A logical reference to an artifact stored via the ArtifactStore abstraction.

    Examples:
    - file:///abs/path/to/artifacts/ab/cdef....json
    - s3://bucket/prefix/ab/cdef....json
    """

    uri: str
    media_type: Optional[str] = None
    bytes: Optional[int] = None

    model_config = {"extra": "allow"}


class ArraySequenceRef(BaseModel):
    """
    Reference to a directory of per-frame arrays.
    """

    format: Literal["npy"] = "npy"
    dir_path: str = Field(..., description="Directory containing per-frame arrays")
    filename_pattern: str = Field("frame_{t:06d}.npy", description="Python format string pattern")
    num_frames: int
    shape_hw: Tuple[int, int]
    dtype: str = "float32"
    units: Units = Units.METERS

    model_config = {"extra": "allow"}


class PoseRef(BaseModel):
    frame_idx: int
    # 4x4 row-major transform (world<-camera), stored as nested lists for JSON.
    T_wc: List[List[float]]
    covariance_uri: Optional[ArtifactURI] = None

    model_config = {"extra": "allow"}


class PoseSet(BaseModel):
    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    poses: List[PoseRef] = Field(default_factory=list)
    stats: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}


class LandmarkRef(BaseModel):
    landmark_id: str
    xyz: Tuple[float, float, float]
    covariance_uri: Optional[ArtifactURI] = None

    model_config = {"extra": "allow"}


class LandmarkSet(BaseModel):
    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    landmarks: List[LandmarkRef] = Field(default_factory=list)
    stats: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}


class TrackObservation(BaseModel):
    frame_idx: int
    xy_px: Tuple[float, float]
    device_id: Optional[str] = None
    confidence: Optional[float] = None

    model_config = {"extra": "allow"}


class TrackRef(BaseModel):
    track_id: str
    observations: List[TrackObservation] = Field(default_factory=list)
    stats: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}


class TrackSet(BaseModel):
    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    tracks: List[TrackRef] = Field(default_factory=list)
    stats: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}


class CalibrationParams(BaseModel):
    """
    Canonical calibration references for a capture bundle/pipeline stage.
    """

    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    intrinsics_by_device: Dict[str, List[List[float]]] = Field(default_factory=dict)
    rig_extrinsics_uri: Optional[ArtifactURI] = None
    sync_offsets_uri: Optional[ArtifactURI] = None

    model_config = {"extra": "allow"}


class Provenance(BaseModel):
    """
    Minimal provenance for auditability.
    """

    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    created_at_unix_s: Optional[float] = None
    git_commit: Optional[str] = None
    config: Dict[str, Any] = Field(default_factory=dict)
    upstream: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}


class MetrologyClaimStatus(str, Enum):
    """
    Whether outputs are allowed to make metrological claims (SPEC §6.7, §11.6).
    """

    METROLOGICAL_OK = "metrological_ok"
    METROLOGICAL_UNKNOWN = "metrological_unknown"
    METROLOGICAL_DISABLED = "metrological_disabled"


class AuditGateOutcome(BaseModel):
    """
    A lightweight copy of audit gate outcomes (SPEC §5.4.3).
    """

    name: str
    passed: bool
    details: Dict[str, Any] = Field(default_factory=dict)

    model_config = {"extra": "allow"}


class TeacherArtifactBundle(BaseModel):
    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    capture_id: str
    device_id: str
    operating_regime: Optional[OperatingRegime] = None
    scene_type: Optional[str] = None
    difficulty_flags: List[str] = Field(default_factory=list)
    metrology_claim: MetrologyClaimStatus = MetrologyClaimStatus.METROLOGICAL_UNKNOWN
    audit_gates: List[AuditGateOutcome] = Field(default_factory=list)
    depth: ArraySequenceRef
    sigma_z: ArraySequenceRef
    calibration: Optional[CalibrationParams] = None
    poses: Optional[PoseSet] = None
    landmarks: Optional[LandmarkSet] = None
    tracks: Optional[TrackSet] = None
    stats: Dict[str, Any] = Field(default_factory=dict)
    provenance: Provenance = Field(default_factory=Provenance)

    model_config = {"extra": "allow"}


class InferenceArtifactBundle(BaseModel):
    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    input: str
    device_id: str
    operating_regime: Optional[OperatingRegime] = None
    scene_type: Optional[str] = None
    difficulty_flags: List[str] = Field(default_factory=list)
    metrology_claim: MetrologyClaimStatus = MetrologyClaimStatus.METROLOGICAL_UNKNOWN
    depth: ArraySequenceRef
    sigma_z: ArraySequenceRef
    eae_uri: Optional[ArtifactURI] = None
    reconstruction_uri: Optional[ArtifactURI] = None
    stats: Dict[str, Any] = Field(default_factory=dict)
    provenance: Provenance = Field(default_factory=Provenance)

    model_config = {"extra": "allow"}


class AuditArtifactBundle(BaseModel):
    schema_version: IntermediateSchemaVersion = IntermediateSchemaVersion.V1_0
    measurements_uri: Optional[ArtifactURI] = None
    result_uri: Optional[ArtifactURI] = None
    calibration_uri: Optional[ArtifactURI] = None
    stats: Dict[str, Any] = Field(default_factory=dict)
    provenance: Provenance = Field(default_factory=Provenance)
    model_config = {"extra": "allow"}