from decoupled_wbc.control.robot_model.robot_model import RobotModel
from decoupled_wbc.data.constants import RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH


def get_modality_config(robot_model: RobotModel, add_stereo_camera: bool = False) -> dict:
    """
    Get the modality config for the robot model.
    """
    left_hand_indices = sorted(robot_model.get_joint_group_indices("left_hand"))
    right_hand_indices = sorted(robot_model.get_joint_group_indices("right_hand"))
    left_arm_indices = sorted(robot_model.get_joint_group_indices("left_arm"))
    right_arm_indices = sorted(robot_model.get_joint_group_indices("right_arm"))
    waist_indices = sorted(robot_model.get_joint_group_indices("waist"))
    left_leg_indices = sorted(robot_model.get_joint_group_indices("left_leg"))
    right_leg_indices = sorted(robot_model.get_joint_group_indices("right_leg"))

    modality_config = {
        "state": {
            "left_leg": {"start": left_leg_indices[0], "end": left_leg_indices[-1] + 1},
            "right_leg": {"start": right_leg_indices[0], "end": right_leg_indices[-1] + 1},
            "waist": {"start": waist_indices[0], "end": waist_indices[-1] + 1},
            "left_arm": {"start": left_arm_indices[0], "end": left_arm_indices[-1] + 1},
            "left_hand": {"start": left_hand_indices[0], "end": left_hand_indices[-1] + 1},
            "right_arm": {"start": right_arm_indices[0], "end": right_arm_indices[-1] + 1},
            "right_hand": {"start": right_hand_indices[0], "end": right_hand_indices[-1] + 1},
            "left_wrist_pos": {"start": 0, "end": 3, "original_key": "observation.eef_state"},
            "left_wrist_abs_quat": {
                "start": 3,
                "end": 7,
                "original_key": "observation.eef_state",
                "rotation_type": "quaternion",
            },
            "right_wrist_pos": {"start": 7, "end": 10, "original_key": "observation.eef_state"},
            "right_wrist_abs_quat": {
                "start": 10,
                "end": 14,
                "original_key": "observation.eef_state",
                "rotation_type": "quaternion",
            },
        },
        "action": {
            "left_leg": {"start": left_leg_indices[0], "end": left_leg_indices[-1] + 1},
            "right_leg": {"start": right_leg_indices[0], "end": right_leg_indices[-1] + 1},
            "waist": {"start": waist_indices[0], "end": waist_indices[-1] + 1},
            "left_arm": {"start": left_arm_indices[0], "end": left_arm_indices[-1] + 1},
            "left_hand": {"start": left_hand_indices[0], "end": left_hand_indices[-1] + 1},
            "right_arm": {"start": right_arm_indices[0], "end": right_arm_indices[-1] + 1},
            "right_hand": {"start": right_hand_indices[0], "end": right_hand_indices[-1] + 1},
            "left_wrist_pos": {"start": 0, "end": 3, "original_key": "action.eef"},
            "left_wrist_abs_quat": {
                "start": 3,
                "end": 7,
                "original_key": "action.eef",
                "rotation_type": "quaternion",
            },
            "right_wrist_pos": {"start": 7, "end": 10, "original_key": "action.eef"},
            "right_wrist_abs_quat": {
                "start": 10,
                "end": 14,
                "original_key": "action.eef",
                "rotation_type": "quaternion",
            },
            "base_height_command": {
                "start": 0,
                "end": 1,
                "original_key": "teleop.base_height_command",
            },
            "navigate_command": {"start": 0, "end": 3, "original_key": "teleop.navigate_command"},
        },
        "video": {"ego_view": {"original_key": "observation.images.ego_view"}},
        "annotation": {"human.task_description": {"original_key": "task_index"}},
    }
    if add_stereo_camera:
        modality_config["video"].update(
            {
                "ego_view_left_mono": {"original_key": "observation.images.ego_view_left_mono"},
                "ego_view_right_mono": {"original_key": "observation.images.ego_view_right_mono"},
            }
        )

    return modality_config


def get_dataset_features(robot_model: RobotModel, add_stereo_camera: bool = False) -> dict:
    """
    Get the dataset features for the robot model.
    """
    dataset_features = {
        "observation.images.ego_view": {
            "dtype": "video",
            "shape": [RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH, 3],
            "names": ["height", "width", "channel"],
        },
        "observation.state": {
            "dtype": "float64",
            "shape": (robot_model.num_joints,),
            "names": robot_model.joint_names,
        },
        "observation.eef_state": {
            "dtype": "float64",
            "shape": (14,),
            "names": [
                "left_wrist_pos",
                "left_wrist_abs_quat",
                "right_wrist_pos",
                "right_wrist_abs_quat",
            ],
        },
        "action": {
            "dtype": "float64",
            "shape": (robot_model.num_joints,),
            "names": robot_model.joint_names,
        },
        "action.eef": {
            "dtype": "float64",
            "shape": (14,),
            "names": [
                "left_wrist_pos",
                "left_wrist_abs_quat",
                "right_wrist_pos",
                "right_wrist_abs_quat",
            ],
        },
        "observation.img_state_delta": {
            "dtype": "float32",
            "shape": (1,),
            "names": "img_state_delta",
        },
        "teleop.navigate_command": {
            "dtype": "float64",
            "shape": (3,),
            "names": ["lin_vel_x", "lin_vel_y", "ang_vel_z"],
        },
        "teleop.base_height_command": {
            "dtype": "float64",
            "shape": (1,),
            "names": "base_height_command",
        },
    }
    if add_stereo_camera:
        dataset_features.update(
            {
                "observation.images.ego_view_left_mono": {
                    "dtype": "video",
                    "shape": [RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH, 3],
                    "names": ["height", "width", "channel"],
                },
                "observation.images.ego_view_right_mono": {
                    "dtype": "video",
                    "shape": [RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH, 3],
                    "names": ["height", "width", "channel"],
                },
            }
        )

    return dataset_features