"""Versioned H2 contracts; G1 checkpoints cannot be loaded into this runtime."""
from __future__ import annotations

from dataclasses import dataclass
import numpy as np

from ..h2_contract import EXPECTED_POLICY_JOINT_NAMES, H2_JOINT_MAP_REVISION

OBS_DIM = 182
BASE_ACTION_DIM = 31
RESIDUAL_ACTION_DIM = 7
ACTOR_INPUT_DIM = OBS_DIM + BASE_ACTION_DIM
FORMAT = "h2-real-residual-td3-v3"
REWARD_PROTOCOL = "h2-gap-return-n5-v1"
RIGHT_ARM_JOINTS = tuple(f"right_{name}_joint" for name in (
    "shoulder_pitch", "shoulder_roll", "shoulder_yaw", "elbow", "wrist_roll", "wrist_pitch", "wrist_yaw"))


def vector(value, size):
    array = np.asarray(value, dtype=np.float32)
    if array.shape != (size,) or not np.isfinite(array).all():
        raise ValueError(f"expected finite vector ({size},), got {array.shape}")
    return array.copy()


@dataclass(frozen=True)
class ActionContract:
    policy_names: tuple[str, ...]
    action_scale: tuple[float, ...]
    residual_scales_rad: tuple[float, ...]

    def __post_init__(self):
        for name in ("policy_names", "action_scale", "residual_scales_rad"):
            object.__setattr__(self, name, tuple(getattr(self, name)))
        for name in ("action_scale", "residual_scales_rad"):
            object.__setattr__(self, name, tuple(float(v) for v in getattr(self, name)))
        if self.policy_names != EXPECTED_POLICY_JOINT_NAMES:
            raise ValueError("residual policy requires the audited H2 policy joint order")
        scales = vector(self.action_scale, BASE_ACTION_DIM)
        residual = vector(self.residual_scales_rad, RESIDUAL_ACTION_DIM)
        if np.any(scales <= 0) or np.any(residual <= 0):
            raise ValueError("action scales must be positive")

    @property
    def indices(self):
        return np.array([self.policy_names.index(n) for n in RIGHT_ARM_JOINTS])

    def to_dict(self):
        return dict(format=FORMAT, robot="h2", obs_dim=OBS_DIM, action_dim=BASE_ACTION_DIM,
                    residual_dim=RESIDUAL_ACTION_DIM, joint_map_revision=H2_JOINT_MAP_REVISION,
                    policy_names=list(self.policy_names), action_scale=list(self.action_scale),
                    residual_scales_rad=list(self.residual_scales_rad),
                    right_arm_joints=list(RIGHT_ARM_JOINTS), action_clip=10.0)

    def validate_saved(self, saved):
        if saved != self.to_dict():
            raise ValueError("checkpoint/action contract mismatch (requires exact H2 dimensions, joint map and scales)")


@dataclass(frozen=True)
class Transition:
    start_s: float
    end_s: float
    base_obs: np.ndarray
    base_action: np.ndarray
    res_action: np.ndarray
    combined_action: np.ndarray
    next_base_obs: np.ndarray
    next_base_action: np.ndarray
    reward: float = 0.0
    terminated: bool = False
    position_distance_m: float | None = None
    position_reward_missing: bool = False

    def __post_init__(self):
        for name, size in (("base_obs", OBS_DIM), ("next_base_obs", OBS_DIM),
                           ("base_action", BASE_ACTION_DIM), ("next_base_action", BASE_ACTION_DIM),
                           ("combined_action", BASE_ACTION_DIM), ("res_action", RESIDUAL_ACTION_DIM)):
            object.__setattr__(self, name, vector(getattr(self, name), size))
        if not np.isfinite([self.start_s, self.end_s, self.reward]).all() or self.end_s <= self.start_s:
            raise ValueError("invalid transition interval/reward")


@dataclass(frozen=True)
class EvaluationTransition(Transition):
    """Measured action interval for scoring, never a residual-training sample."""
    observation_dim: int = OBS_DIM

    def __post_init__(self):
        if not isinstance(self.observation_dim, int) or self.observation_dim <= 0:
            raise ValueError("evaluation observation_dim must be a positive integer")
        for name, size in (("base_obs", self.observation_dim), ("next_base_obs", self.observation_dim),
                           ("base_action", BASE_ACTION_DIM), ("next_base_action", BASE_ACTION_DIM),
                           ("combined_action", BASE_ACTION_DIM), ("res_action", RESIDUAL_ACTION_DIM)):
            object.__setattr__(self, name, vector(getattr(self, name), size))
        if not np.isfinite([self.start_s, self.end_s, self.reward]).all() or self.end_s <= self.start_s:
            raise ValueError("invalid transition interval/reward")


@dataclass(frozen=True)
class EpisodeBatch:
    sequence: int
    command_id: int
    phase: str
    policy_version: int
    transitions: tuple[Transition, ...]
    summary: dict
    collector_rng_state: dict

    def validate(self):
        if self.sequence <= 0 or self.command_id <= 0 or self.phase not in ("warmup", "online"):
            raise ValueError("invalid episode identity/phase")
        if not self.transitions or not self.transitions[-1].terminated:
            raise ValueError("only complete, adjudicated episodes may enter the learner")
        for i, item in enumerate(self.transitions):
            if i < len(self.transitions) - 1 and item.terminated:
                raise ValueError("termination must occur exactly once at the end")
            if i and abs(self.transitions[i-1].end_s - item.start_s) > 1e-6:
                raise ValueError("episode transitions must be contiguous")