"""Hugging Face configuration for MiniCPM-RobotTrack."""

from typing import Any, Dict, Optional

from transformers.configuration_utils import PretrainedConfig


class MiniCPMRobotTrackConfig(PretrainedConfig):
    """Configuration for the MiniCPM-RobotTrack trajectory policy."""

    model_type = "minicpm_robottrack"
    has_no_defaults_at_init = True

    def __init__(
        self,
        backbone_config: Optional[Dict[str, Any]] = None,
        vision_feature_dim: int = 1536,
        history_frames: int = 31,
        coarse_tokens_per_frame: int = 4,
        fine_tokens_current_frame: int = 64,
        num_waypoints: int = 8,
        action_dim: int = 3,
        max_text_tokens: int = 128,
        max_time_steps: int = 4096,
        trajectory_dropout: float = 0.4,
        xy_scale: float = 2.0,
        use_tanh_actions: bool = True,
        backbone_dtype: str = "bfloat16",
        **kwargs: Any,
    ) -> None:
        self.backbone_config = dict(backbone_config or {})
        self.vision_feature_dim = int(vision_feature_dim)
        self.history_frames = int(history_frames)
        self.coarse_tokens_per_frame = int(coarse_tokens_per_frame)
        self.fine_tokens_current_frame = int(fine_tokens_current_frame)
        self.num_waypoints = int(num_waypoints)
        self.action_dim = int(action_dim)
        self.max_text_tokens = int(max_text_tokens)
        self.max_time_steps = int(max_time_steps)
        self.trajectory_dropout = float(trajectory_dropout)
        self.xy_scale = float(xy_scale)
        self.use_tanh_actions = bool(use_tanh_actions)
        self.backbone_dtype = str(backbone_dtype)
        self._validate_robot_track_config()
        super().__init__(**kwargs)

    def _validate_robot_track_config(self) -> None:
        if not self.backbone_config:
            raise ValueError("backbone_config must contain the bundled MiniCPM configuration")
        if self.vision_feature_dim <= 0:
            raise ValueError("vision_feature_dim must be positive")
        if self.history_frames < 0:
            raise ValueError("history_frames cannot be negative")
        for name, value in (
            ("coarse_tokens_per_frame", self.coarse_tokens_per_frame),
            ("fine_tokens_current_frame", self.fine_tokens_current_frame),
        ):
            side = int(round(value**0.5)) if value > 0 else 0
            if side * side != value:
                raise ValueError(f"{name} must be a positive square number")
        if self.num_waypoints < 2:
            raise ValueError("num_waypoints must be at least 2")
        if self.action_dim != 3:
            raise ValueError("actions must use [x, y, yaw] format")
        if self.max_text_tokens <= 0 or self.max_time_steps <= 0:
            raise ValueError("token and time limits must be positive")
        if not 0.0 <= self.trajectory_dropout < 1.0:
            raise ValueError("trajectory_dropout must be in [0, 1)")
        if self.xy_scale <= 0.0:
            raise ValueError("xy_scale must be positive")