dovla_cil/data/schema.py

from __future__ import annotations

import hashlib
import json
import math
from dataclasses import asdict, dataclass, field, replace
from pathlib import Path
from typing import Any, Iterable, Iterator

JSONValue = str | int | float | bool | None | list["JSONValue"] | dict[str, "JSONValue"]
Observation = dict[str, JSONValue]
CIL_VERSION = "0.1"


def assert_jsonable(value: Any, field_name: str) -> None:
    try:
        json.dumps(value)
    except TypeError as exc:
        raise TypeError(f"{field_name} must be JSON serializable") from exc


@dataclass(frozen=True)
class ActionChunk:
    action_id: str = ""
    representation: str = "numeric"
    horizon: int = 1
    values: list[list[float]] | dict[str, JSONValue] | list[dict[str, JSONValue]] = field(
        default_factory=list
    )
    skill_type: str | None = None
    metadata: dict[str, JSONValue] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.horizon <= 0:
            raise ValueError("ActionChunk.horizon must be positive")
        if not self.representation:
            raise ValueError("ActionChunk.representation must be non-empty")
        values = _normalize_action_values(self.values)
        assert_jsonable(values, "ActionChunk.values")
        assert_jsonable(self.metadata, "ActionChunk.metadata")
        action_id = self.action_id or "act-" + _stable_json_hash(
            {
                "representation": self.representation,
                "horizon": self.horizon,
                "values": values,
                "skill_type": self.skill_type,
                "metadata": self.metadata,
            },
            length=16,
        )
        object.__setattr__(self, "values", values)
        object.__setattr__(self, "action_id", action_id)

    @property
    def flat_values(self) -> list[float]:
        if isinstance(self.values, list) and all(isinstance(row, list) for row in self.values):
            return [float(item) for row in self.values for item in row]
        return []

    def to_dict(self) -> dict[str, JSONValue]:
        return asdict(self)

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "ActionChunk":
        return cls(
            action_id=str(payload.get("action_id", "")),
            representation=str(payload.get("representation", "numeric")),
            horizon=int(payload.get("horizon", 1)),
            values=payload.get("values", []),
            skill_type=payload.get("skill_type"),
            metadata=dict(payload.get("metadata", {})),
        )


@dataclass(frozen=True)
class StructuredEffect:
    object_pose_delta: dict[str, list[float]] = field(default_factory=dict)
    contact_events: list[dict[str, JSONValue]] = field(default_factory=list)
    relation_before: dict[str, bool] = field(default_factory=dict)
    relation_after: dict[str, bool] = field(default_factory=dict)
    grasp_success: bool | None = None
    moved_objects: list[str] = field(default_factory=list)
    articulation_delta: dict[str, float] = field(default_factory=dict)
    symbolic_before: dict[str, JSONValue] = field(default_factory=dict)
    symbolic_after: dict[str, JSONValue] = field(default_factory=dict)
    metadata: dict[str, JSONValue] = field(default_factory=dict)

    def __post_init__(self) -> None:
        pose_delta = {
            str(key): [float(value) for value in values]
            for key, values in self.object_pose_delta.items()
        }
        articulation_delta = {
            str(key): float(value) for key, value in self.articulation_delta.items()
        }
        if any(not math.isfinite(value) for values in pose_delta.values() for value in values):
            raise ValueError("StructuredEffect.object_pose_delta values must be finite")
        if any(not math.isfinite(value) for value in articulation_delta.values()):
            raise ValueError("StructuredEffect.articulation_delta values must be finite")
        object.__setattr__(self, "object_pose_delta", pose_delta)
        object.__setattr__(self, "articulation_delta", articulation_delta)
        object.__setattr__(
            self,
            "relation_before",
            {str(key): bool(value) for key, value in self.relation_before.items()},
        )
        object.__setattr__(
            self,
            "relation_after",
            {str(key): bool(value) for key, value in self.relation_after.items()},
        )
        object.__setattr__(self, "moved_objects", [str(item) for item in self.moved_objects])
        assert_jsonable(self.contact_events, "StructuredEffect.contact_events")
        assert_jsonable(self.symbolic_before, "StructuredEffect.symbolic_before")
        assert_jsonable(self.symbolic_after, "StructuredEffect.symbolic_after")
        assert_jsonable(self.metadata, "StructuredEffect.metadata")

    @property
    def metrics(self) -> dict[str, float]:
        """Compatibility view for older effect smoke tests."""

        metrics = {
            f"pose_delta_norm/{object_id}": math.sqrt(sum(value * value for value in delta))
            for object_id, delta in self.object_pose_delta.items()
        }
        metrics.update(self.articulation_delta)
        metrics.update(
            {
                str(key): float(value)
                for key, value in self.metadata.get("metrics", {}).items()
                if isinstance(value, (int, float))
            }
        )
        return metrics

    @property
    def predicates(self) -> dict[str, bool]:
        predicates = dict(self.relation_after)
        predicates.update(
            {
                str(key): bool(value)
                for key, value in self.metadata.get("predicates", {}).items()
                if isinstance(value, bool)
            }
        )
        return predicates

    def to_dict(self) -> dict[str, JSONValue]:
        return asdict(self)

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "StructuredEffect":
        return cls(
            object_pose_delta=dict(payload.get("object_pose_delta", {})),
            contact_events=list(payload.get("contact_events", [])),
            relation_before=dict(payload.get("relation_before", {})),
            relation_after=dict(payload.get("relation_after", {})),
            grasp_success=payload.get("grasp_success"),
            moved_objects=list(payload.get("moved_objects", [])),
            articulation_delta=dict(payload.get("articulation_delta", {})),
            symbolic_before=dict(payload.get("symbolic_before", {})),
            symbolic_after=dict(payload.get("symbolic_after", {})),
            metadata=dict(payload.get("metadata", {})),
        )


@dataclass(frozen=True)
class RewardInfo:
    progress: float
    success: bool
    terminal_success: bool
    dense_components: dict[str, float] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not math.isfinite(float(self.progress)):
            raise ValueError("RewardInfo.progress must be finite")
        components = {str(key): float(value) for key, value in self.dense_components.items()}
        if any(not math.isfinite(value) for value in components.values()):
            raise ValueError("RewardInfo.dense_components values must be finite")
        object.__setattr__(self, "progress", float(self.progress))
        object.__setattr__(self, "success", bool(self.success))
        object.__setattr__(self, "terminal_success", bool(self.terminal_success))
        object.__setattr__(self, "dense_components", components)

    @property
    def score(self) -> float:
        return float(self.progress) + (1.0 if self.terminal_success else 0.0)

    def to_dict(self) -> dict[str, JSONValue]:
        return asdict(self)

    @classmethod
    def from_dict(cls, payload: dict[str, Any] | float | int) -> "RewardInfo":
        if isinstance(payload, (float, int)):
            progress = float(payload)
            return cls(progress=progress, success=progress > 0.0, terminal_success=progress > 0.0)
        return cls(
            progress=float(payload["progress"]),
            success=bool(payload["success"]),
            terminal_success=bool(payload["terminal_success"]),
            dense_components=dict(payload.get("dense_components", {})),
        )


@dataclass(frozen=True)
class FailureInfo:
    type: str
    symbolic_reason: str | None = None
    language_explanation: str | None = None
    metadata: dict[str, JSONValue] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.type:
            raise ValueError("FailureInfo.type must be non-empty")
        assert_jsonable(self.metadata, "FailureInfo.metadata")

    def to_dict(self) -> dict[str, JSONValue]:
        return asdict(self)

    @classmethod
    def from_dict(cls, payload: dict[str, Any] | None) -> "FailureInfo | None":
        if payload is None:
            return None
        return cls(
            type=str(payload["type"]),
            symbolic_reason=payload.get("symbolic_reason"),
            language_explanation=payload.get("language_explanation"),
            metadata=dict(payload.get("metadata", {})),
        )


@dataclass(frozen=True)
class CILRecord:
    version: str
    record_id: str
    group_id: str
    state_hash: str
    task_id: str
    scene_id: str | None
    instruction: str
    instruction_family: dict[str, JSONValue]
    observation_ref: str | None
    observation_inline: dict[str, JSONValue] | None
    action_chunk: ActionChunk
    next_observation_ref: str | None
    next_observation_inline: dict[str, JSONValue] | None
    structured_effect: StructuredEffect
    reward: RewardInfo
    regret: float | None
    rank_within_group: int | None
    candidate_type: str
    failure: FailureInfo | None
    metadata: dict[str, JSONValue] = field(default_factory=dict)

    def __post_init__(self) -> None:
        required_fields = (
            "version",
            "record_id",
            "group_id",
            "state_hash",
            "task_id",
            "instruction",
        )
        for field_name in required_fields:
            if not getattr(self, field_name):
                raise ValueError(f"CILRecord.{field_name} must be non-empty")
        if not self.observation_ref and self.observation_inline is None:
            raise ValueError("CILRecord requires observation_ref or observation_inline")
        if not self.next_observation_ref and self.next_observation_inline is None:
            raise ValueError("CILRecord requires next_observation_ref or next_observation_inline")
        regret = None if self.regret is None else float(self.regret)
        if regret is not None and not math.isfinite(regret):
            raise ValueError("CILRecord.regret must be finite")
        if self.rank_within_group is not None and self.rank_within_group < 0:
            raise ValueError("CILRecord.rank_within_group must be non-negative")
        if not self.candidate_type:
            raise ValueError("CILRecord.candidate_type must be non-empty")
        object.__setattr__(self, "regret", regret)
        if isinstance(self.action_chunk, dict):
            object.__setattr__(self, "action_chunk", ActionChunk.from_dict(self.action_chunk))
        if isinstance(self.structured_effect, dict):
            object.__setattr__(
                self, "structured_effect", StructuredEffect.from_dict(self.structured_effect)
            )
        if isinstance(self.reward, (dict, float, int)):
            object.__setattr__(self, "reward", RewardInfo.from_dict(self.reward))
        if isinstance(self.failure, dict):
            object.__setattr__(self, "failure", FailureInfo.from_dict(self.failure))
        assert_jsonable(self.instruction_family, "CILRecord.instruction_family")
        assert_jsonable(self.observation_inline, "CILRecord.observation_inline")
        assert_jsonable(self.next_observation_inline, "CILRecord.next_observation_inline")
        assert_jsonable(self.metadata, "CILRecord.metadata")

    @property
    def action(self) -> ActionChunk:
        return self.action_chunk

    @property
    def next_observation(self) -> dict[str, JSONValue] | None:
        return self.next_observation_inline

    @property
    def reward_value(self) -> float:
        return self.reward.score

    def validate(self) -> None:
        assert_jsonable(self.to_dict(), "CILRecord")

    def to_dict(self) -> dict[str, JSONValue]:
        payload = asdict(self)
        payload["action_chunk"] = self.action_chunk.to_dict()
        payload["structured_effect"] = self.structured_effect.to_dict()
        payload["reward"] = self.reward.to_dict()
        payload["failure"] = self.failure.to_dict() if self.failure else None
        return payload

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "CILRecord":
        return cls(
            version=str(payload.get("version", CIL_VERSION)),
            record_id=str(payload["record_id"]),
            group_id=str(payload["group_id"]),
            state_hash=str(payload["state_hash"]),
            task_id=str(payload["task_id"]),
            scene_id=payload.get("scene_id"),
            instruction=str(payload["instruction"]),
            instruction_family=dict(payload.get("instruction_family", {})),
            observation_ref=payload.get("observation_ref"),
            observation_inline=payload.get("observation_inline"),
            action_chunk=ActionChunk.from_dict(dict(payload["action_chunk"])),
            next_observation_ref=payload.get("next_observation_ref"),
            next_observation_inline=payload.get("next_observation_inline"),
            structured_effect=StructuredEffect.from_dict(dict(payload["structured_effect"])),
            reward=RewardInfo.from_dict(payload["reward"]),
            regret=payload.get("regret"),
            rank_within_group=payload.get("rank_within_group"),
            candidate_type=str(payload["candidate_type"]),
            failure=FailureInfo.from_dict(payload.get("failure")),
            metadata=dict(payload.get("metadata", {})),
        )


@dataclass(frozen=True)
class CILGroup:
    group_id: str
    state_hash: str
    task_id: str
    instruction: str
    records: list[CILRecord]

    def __post_init__(self) -> None:
        validate_group(self.records)
        if self.records:
            first = self.records[0]
            if (
                self.group_id != first.group_id
                or self.state_hash != first.state_hash
                or self.task_id != first.task_id
                or self.instruction != first.instruction
            ):
                raise ValueError("CILGroup metadata does not match contained records")

    def to_dict(self) -> dict[str, JSONValue]:
        return {
            "group_id": self.group_id,
            "state_hash": self.state_hash,
            "task_id": self.task_id,
            "instruction": self.instruction,
            "records": [record.to_dict() for record in self.records],
        }

    @classmethod
    def from_records(cls, records: list[CILRecord]) -> "CILGroup":
        validate_group(records)
        first = records[0]
        return cls(
            group_id=first.group_id,
            state_hash=first.state_hash,
            task_id=first.task_id,
            instruction=first.instruction,
            records=list(records),
        )


def make_record_id(group_id: str, action_id: str, seed: int) -> str:
    if not group_id:
        raise ValueError("group_id must be non-empty")
    if not action_id:
        raise ValueError("action_id must be non-empty")
    digest = hashlib.sha256(f"{group_id}:{action_id}:{seed}".encode("utf-8")).hexdigest()
    return f"rec-{digest[:24]}"


def compute_state_hash(state_blob: bytes) -> str:
    return hashlib.sha256(state_blob).hexdigest()


def compute_regret_and_ranks(records: list[CILRecord]) -> list[CILRecord]:
    if not records:
        return []
    validate_group(records)
    scored = [(index, _reward_score(record)) for index, record in enumerate(records)]
    best_score = max(score for _, score in scored)
    ordered_indices = [
        index
        for index, _score in sorted(
            scored, key=lambda item: (-item[1], records[item[0]].record_id)
        )
    ]
    ranks = {index: rank for rank, index in enumerate(ordered_indices)}
    return [
        replace(
            record,
            regret=best_score - _reward_score(record),
            rank_within_group=ranks[index],
        )
        for index, record in enumerate(records)
    ]


def validate_group(records: list[CILRecord]) -> None:
    if not records:
        raise ValueError("CIL group must contain at least one record")
    group_ids = {record.group_id for record in records}
    state_hashes = {record.state_hash for record in records}
    task_ids = {record.task_id for record in records}
    if len(group_ids) != 1:
        raise ValueError("All CIL records in a group must share group_id")
    if len(state_hashes) != 1:
        raise ValueError("All CIL records in a group must share state_hash")
    if len(task_ids) != 1:
        raise ValueError("All CIL records in a group must share task_id")


def write_cil_jsonl(records: Iterable[CILRecord], path: str | Path) -> None:
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    with target.open("w", encoding="utf-8") as handle:
        for record in records:
            record.validate()
            handle.write(json.dumps(record.to_dict(), sort_keys=True) + "\n")


def iter_cil_jsonl(path: str | Path) -> Iterator[CILRecord]:
    with Path(path).open("r", encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                yield CILRecord.from_dict(json.loads(line))


def _reward_score(record: CILRecord) -> float:
    return record.reward.score


def _stable_json_hash(payload: Any, *, length: int) -> str:
    encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode(
        "utf-8"
    )
    return hashlib.sha256(encoded).hexdigest()[:length]


def _normalize_action_values(
    values: Any,
) -> list[list[float]] | dict[str, JSONValue] | list[dict[str, JSONValue]]:
    if isinstance(values, tuple):
        values = list(values)
    if isinstance(values, dict):
        return dict(values)
    if isinstance(values, list) and not values:
        return values
    if isinstance(values, list) and all(isinstance(item, dict) for item in values):
        return [dict(item) for item in values]
    if isinstance(values, list) and all(isinstance(item, (int, float)) for item in values):
        return [[float(item) for item in values]]
    if isinstance(values, list) and all(isinstance(row, (list, tuple)) for row in values):
        normalized: list[list[float]] = []
        for row in values:
            row_values = [float(item) for item in row]
            if any(not math.isfinite(item) for item in row_values):
                raise ValueError("ActionChunk.values numeric entries must be finite")
            normalized.append(row_values)
        return normalized
    raise TypeError("ActionChunk.values must be numeric rows, a dict, or a list of dicts")