# Copyright 2025 The RLinf Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import asdict, dataclass, field

import numpy as np


@dataclass
class Turtle2RobotState:
    """Turtle2 robot state including followers, head, lift, and car pose.

    Attributes:
        follow1_pos: Follower 1 position (7-dim).
        follow1_joints: Follower 1 joint angles (7-dim).
        follow1_cur_data: Follower 1 current data (7-dim).
        follow2_pos: Follower 2 position (7-dim).
        follow2_joints: Follower 2 joint angles (7-dim).
        follow2_cur_data: Follower 2 current data (7-dim).
        head_pos: Head position (2-dim).
        lift: Lift height.
        car_pose: Car pose [x, y, theta] (3-dim).
    """

    follow1_pos: np.ndarray = field(default_factory=lambda: np.zeros(7))
    follow1_joints: np.ndarray = field(default_factory=lambda: np.zeros(7))
    follow1_cur_data: np.ndarray = field(default_factory=lambda: np.zeros(7))
    follow2_pos: np.ndarray = field(default_factory=lambda: np.zeros(7))
    follow2_joints: np.ndarray = field(default_factory=lambda: np.zeros(7))
    follow2_cur_data: np.ndarray = field(default_factory=lambda: np.zeros(7))

    head_pos: np.ndarray = field(default_factory=lambda: np.zeros(2))
    lift: float = 0.0
    car_pose: np.ndarray = field(default_factory=lambda: np.zeros(3))

    def to_dict(self):
        """Convert the dataclass to a serializable dictionary."""
        return asdict(self)
