import torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset

dataset = LeRobotDataset("RoboCOIN/R1_Lite_open_and_close_microwave_oven")

dataloader = torch.utils.data.DataLoader(
    dataset,
    num_workers=8,
    batch_size=32,
)

# Base configuration class for joint control
@RobotConfig.register_subclass("base_robot")
@dataclass
class BaseRobotConfig(RobotConfig):
    # Camera settings, represented as dictionary, key is camera name, value is camera config class, e.g.
    # {
    #     head: {type: opencv, index_or_path:0, height: 480, width: 640, fps: 30}, 
    #     wrist: {type: opencv, index_or_path:1, height: 480, width: 640, fps: 30},
    # }
    # The above example creates head and wrist cameras, loading /dev/video0, /dev/video1 respectively
    # Finally sent to model: {"observation.head": shape(480, 640, 3), "observation.wrist": shape(480, 640, 3)}
    cameras: dict[str, CameraConfig] = field(default_factory=dict)
    # Joint names, including gripper
    joint_names: list[str] = field(default_factory=lambda: [
        'joint_1', 'joint_2', 'joint_3', 'joint_4', 'joint_5', 'joint_6', 'joint_7', 'gripper',
    ]) 

    # Initialization mode: none for no initialization, joint/end_effector for joint/end effector based initialization
    init_type: str = 'none'
    # Values to initialize before starting inference based on initialization mode
    # For joint: units in radian
    # For end_effector: units in m (first 3 values) / radian (values 3~6)
    init_state: list[float] = field(default_factory=lambda: [
        0, 0, 0, 0, 0, 0, 0, 0,
    ])

    # Joint control units, depends on SDK, e.g. Realman SDK has 7 joints receiving angles as parameters, should set:
    # ['degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'm']
    # Last dimension is m, meaning gripper value doesn't need unit conversion
    joint_units: list[str] = field(default_factory=lambda: [
        'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'm',
    ])
    # End effector control units, depends on SDK, e.g. Realman SDK receives meters for xyz and degrees for rpy, should set:
    # ['m', 'm', 'm', 'degree', 'degree', 'degree', 'm']
    # Last dimension is m, meaning gripper value doesn't need unit conversion
    pose_units: list[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'radian', 'radian', 'radian', 'm',
    ])
    # Model input joint control units, depends on dataset, e.g. if dataset saves in radians, should set:
    # ['radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'm']
    # Last dimension is m, meaning gripper value doesn't need unit conversion
    model_joint_units: list[str] = field(default_factory=lambda: [
        'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'm',
    ])
    
    # Relative position control mode: none for absolute position control, previous/init for relative transformation based on previous/initial state
    # Taking joint control as example:
    # - If previous: obtained state + previous state -> target state
    # - If init: obtained state + initial state -> target state
    delta_with: str = 'none'

    # Whether to enable visualization
    visualize: bool = True
    # Whether to draw 2D trajectory, including end effector trajectory on XY, XZ, YZ planes
    draw_2d: bool = True
    # Whether to draw 3D trajectory
    draw_3d: bool = True


# Base configuration class for end effector control
@RobotConfig.register_subclass("base_robot_end_effector")
@dataclass
class BaseRobotEndEffectorConfig(BaseRobotConfig):
    # Relative transformation angles, applicable for cross-body scenarios where different bodies have different zero pose orientations
    base_euler: list[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])

    # Model input end effector control units, depends on dataset, e.g. if dataset saves in meters and radians, should set:
    # ['m', 'm', 'm', 'radian', 'radian', 'radian', 'm']
    # Last dimension is m, meaning gripper value doesn't need unit conversion
    model_pose_units: list[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'radian', 'radian', 'radian', 'm',
    ])

# Dual-arm robot configuration
@RobotConfig.register_subclass("bi_base_robot")
@dataclass
class BiBaseRobotConfig(BaseRobotConfig):
    # Left arm initial pose
    init_state_left: List[float] = field(default_factory=lambda: [
        0, 0, 0, 0, 0, 0, 0, 0,
    ])
    # Right arm initial pose
    init_state_right: List[float] = field(default_factory=lambda: [
        0, 0, 0, 0, 0, 0, 0, 0,
    ])


# Dual-arm robot end effector configuration
@RobotConfig.register_subclass("bi_base_robot_end_effector")
@dataclass
class BiBaseRobotEndEffectorConfig(BiBaseRobotConfig, BaseRobotEndEffectorConfig):
    pass

@RobotConfig.register_subclass("realman")
@dataclass
class RealmanConfig(BaseRobotConfig):
    ip: str = "169.254.128.18" # Realman SDK connection IP
    port: int = 8080           # Realman SDK connection port
    block: bool = False        # Whether to use blocking control
    wait_second: float = 0.1   # If non-blocking, delay after each action
    velocity: int = 30         # Movement velocity

    # Realman has 7 joints + gripper
    joint_names: list[str] = field(default_factory=lambda: [
        'joint_1', 'joint_2', 'joint_3', 'joint_4', 'joint_5', 'joint_6', 'joint_7', 'gripper',
    ])

    # Use joint control to reach Realman's initial task pose
    init_type: str = "joint"
    init_state: list[float] = field(default_factory=lambda: [
        -0.84, -2.03,  1.15,  1.15,  2.71,  1.60, -2.99, 888.00,
    ])

    # Realman SDK defaults to meters + degrees
    joint_units: list[str] = field(default_factory=lambda: [
        'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'm',
    ])
    pose_units: list[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'degree', 'degree', 'degree', 'm',
    ])


@RobotConfig.register_subclass("realman_end_effector")
@dataclass
class RealmanEndEffectorConfig(RealmanConfig, BaseRobotEndEffectorConfig):
    pass

# Dual-arm Realman configuration
@RobotConfig.register_subclass("bi_realman")
@dataclass
class BiRealmanConfig(BiBaseRobotConfig):
    ip_left: str = "169.254.128.18" # Realman left arm SDK connection IP
    port_left: int = 8080 # Realman left arm SDK connection port
    ip_right: str = "169.254.128.19" # Realman right arm SDK connection IP
    port_right: int = 8080 # Realman right arm SDK connection port
    block: bool = False # Whether to use blocking control
    wait_second: float = 0.1 # If non-blocking, delay after each action
    velocity: int = 30 # Movement velocity
    
    # Realman has 7 joints + gripper
    joint_names: List[str] = field(default_factory=lambda: [
        'joint_1', 'joint_2', 'joint_3', 'joint_4', 'joint_5', 'joint_6', 'joint_7', 'gripper',
    ])
    
    # Use joint control to reach Realman's initial task pose
    init_type: str = "joint"
    init_state_left: List[float] = field(default_factory=lambda: [
        -0.84, -2.03,  1.15,  1.15,  2.71,  1.60, -2.99, 888.00,
    ])
    init_state_right: List[float] = field(default_factory=lambda: [
         1.16,  2.01, -0.79, -0.68, -2.84, -1.61,  2.37, 832.00,
    ])

    # Realman SDK defaults to meters + degrees
    joint_units: List[str] = field(default_factory=lambda: [
        'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'm',
    ])
    pose_units: List[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'degree', 'degree', 'degree', 'm',
    ])


# Dual-arm Realman end effector configuration
@RobotConfig.register_subclass("bi_realman_end_effector")
@dataclass
class BiRealmanEndEffectorConfig(BiRealmanConfig, BiBaseRobotEndEffectorConfig):
    pass

@dataclass
class TaskConfig:
    # Scene description
    scene: str = "a yellow basket and a grey towel are place on a white table, the basket is on the left and the towel is on the right."
    # Task instruction
    task: str = "put the towel into the basket."
    # Subtask instructions
    subtasks: List[str] = field(default_factory=lambda: [
        "left gripper catch basket",
        "left gripper move basket to center",
        "right gripper catch towel",
        "right gripper move towel over basket and release",
        "end",
    ])
    # State statistics operators
    operaters: List[Dict] = field(default_factory=lambda: [
        {
            'type': 'position',
            'name': 'position_left',
            'window_size': 1,
            'state_key': 'observation.state',
            'xyz_range': (0, 3),
        }, {
            'type': 'position',
            'name': 'position_right',
            'window_size': 1,
            'state_key': 'observation.state',
            'xyz_range': (7, 10),
        }, {
            'type': 'position_rotation',
            'name': 'position_aligned_left',
            'window_size': 1,
            'position_key': 'position_left',
            'rotation_euler': (0, 0, 0.5 * math.pi),
        }, {
            'type': 'position_rotation',
            'name': 'position_aligned_right',
            'window_size': 1,
            'position_key': 'position_right',
            'rotation_euler': (0, 0, 0.5 * math.pi),
        }, {
            'type': 'movement',
            'name': 'movement_left',
            'window_size': 3,
            'position_key': 'position_aligned_left',
        }, {
            'type': 'movement',
            'name': 'movement_right',
            'window_size': 3,
            'position_key': 'position_aligned_right',
        },{
            'type': 'movement_summary',
            'name': 'movement_summary_left',
            'movement_key': 'movement_left',
            'threshold': 2e-3,
        }, {
            'type': 'movement_summary',
            'name': 'movement_summary_right',
            'movement_key': 'movement_right',
            'threshold': 2e-3,
        }, 
    ])

"""
Configuration for BaseLeader teleoperator.
"""

from dataclasses import dataclass, field
from typing import List

from ..config import TeleoperatorConfig


@TeleoperatorConfig.register_subclass("base_leader")
@dataclass
class BaseLeaderConfig(TeleoperatorConfig):
    """
    Configuration for BaseLeader teleoperator.
    Params:
    - joint_names: List[str], list of joint names, including gripper
    - joint_units: List[str], units for robot joints, for sdk control
    - pose_units: List[str], units for end effector pose, for sdk control
    - model_joint_units: List[str], units for model joints, for model input/output
    """

    # list of joint names, including gripper
    joint_names: List[str] = field(default_factory=lambda: [
        'joint_1', 'joint_2', 'joint_3', 'joint_4', 'joint_5', 'joint_6', 'joint_7', 'gripper',
    ])

    # initialization type and state, each episode start with this state
    # 'none': no initialization
    # 'joint': initialize joint positions
    # 'end_effector': initialize end effector pose
    init_type: str = "none"  # 'none', 'joint', 'end_effector'
    init_state: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    init_threshold: float = 0.1  # threshold to consider initialized

    # units for robot joints, for sdk control
    joint_units: List[str] = field(default_factory=lambda: [
        'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'm',
    ])
    # units for end effector pose, for sdk control
    pose_units: List[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'radian', 'radian', 'radian', 'm',
    ])
    # units for model joints, for model input/output
    model_joint_units: List[str] = field(default_factory=lambda: [
        'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'radian', 'm',
    ])


@TeleoperatorConfig.register_subclass("base_leader_end_effector")
@dataclass
class BaseLeaderEndEffectorConfig(BaseLeaderConfig):
    """
    Configuration for BaseLeaderEndEffector teleoperator.
    Extends BaseLeaderConfig with end-effector specific parameters.
    Params:
    - base_euler: List[float], robot SDK control coordinate system rotation 
      relative to the model coordinate system (not implemented yet)
    - model_pose_units: List[str], units for model end effector pose for model input/output
    """

    # robot SDK control coordinate system rotation relative to the model coordinate system
    base_euler: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
    # units for model end effector pose, for model input/output
    model_pose_units: List[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'radian', 'radian', 'radian', 'm',
    ])

"""
Configuration for Bi-Base Leader Teleoperator
"""

from dataclasses import dataclass

from lerobot.teleoperators import TeleoperatorConfig

from ..base_leader import BaseLeaderConfig, BaseLeaderEndEffectorConfig


@TeleoperatorConfig.register_subclass("bi_base_leader")
@dataclass
class BiBaseLeaderConfig(BaseLeaderConfig):
    """
    Configuration for Bi-Base Leader Teleoperator with joint control
    """

    pass


@TeleoperatorConfig.register_subclass("bi_base_leader_end_effector")
@dataclass
class BiBaseLeaderEndEffectorConfig(BiBaseLeaderConfig, BaseLeaderEndEffectorConfig):
    """
    Configuration for Bi-Base Leader Teleoperator with end effector control
    """
    
    pass

"""
Configuration for Realman leader teleoperator.
"""

from dataclasses import dataclass, field
from typing import List

from ..base_leader import BaseLeaderConfig, BaseLeaderEndEffectorConfig
from ..config import TeleoperatorConfig


@TeleoperatorConfig.register_subclass("realman_leader")
@dataclass
class RealmanLeaderConfig(BaseLeaderConfig):
    """
    Configuration for Realman leader.
    Params:
    - ip: str, IP address of the Realman robot controller
    - port: int, port number for the Realman robot controller
    - block: bool, if True, SDK commands will block until the action is completed
    - wait_second: float, time to wait for non-blocking commands
    - velocity: int, default velocity for joint movements (0-100)
    - joint_names: List[str], list of joint names for the robot, including gripper
    - init_type: str, initialization type, choices: 'none', 'joint', 'end_effector'
    - init_state: List[float], initial joint state for the Realman leader
    - joint_units: List[str], units for robot joints, for sdk control
    - pose_units: List[str], units for end effector pose, for sdk control
    """

    ##### Realman SDK settings #####
    # IP and port of the Realman robot controller
    ip: str = "169.254.128.18"
    port: int = 8080

    # Realman has 7 joints + gripper
    joint_names: List[str] = field(default_factory=lambda: [
        'joint_1', 'joint_2', 'joint_3', 'joint_4', 'joint_5', 'joint_6', 'joint_7', 'gripper',
    ])
    
    # Default initial state for the Realman leader
    init_type: str = "joint"
    init_state: List[float] = field(default_factory=lambda: [
        -0.84, -2.03,  1.15,  1.15,  2.71,  1.60, -2.99, 888.00,
    ])

    # Realman SDK uses degrees for joint angles and meters for positions
    joint_units: List[str] = field(default_factory=lambda: [
        'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'm',
    ])
    pose_units: List[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'degree', 'degree', 'degree', 'm',
    ])


@TeleoperatorConfig.register_subclass("realman_leader_end_effector")
@dataclass
class RealmanLeaderEndEffectorConfig(RealmanLeaderConfig, BaseLeaderEndEffectorConfig):
    """
    Configuration for Realman leader with end effector.
    """

    pass

"""
Configuration for Bi-Realman leader.
"""

from dataclasses import dataclass, field
from typing import List

from ..base_leader import BaseLeaderConfig, BaseLeaderEndEffectorConfig
from ..config import TeleoperatorConfig


@TeleoperatorConfig.register_subclass("bi_realman_leader")
@dataclass
class BiRealmanLeaderConfig(BaseLeaderConfig):
    """
    Configuration for Bi-Realman leader.
    """
    ##### Realman SDK settings #####
    # IP and port settings for the left and right Realman robots.
    ip_left: str = "169.254.128.18"
    port_left: int = 8080
    ip_right: str = "169.254.128.19"
    port_right: int = 8080

    # BiRealman robot has 7 joints and a gripper for each side
    joint_names: List[str] = field(default_factory=lambda: [
        'joint_1', 'joint_2', 'joint_3', 'joint_4', 'joint_5', 'joint_6', 'joint_7', 'gripper',
    ])

    # Default initial state for the BiRealman leader
    init_type: str = "joint"
    init_state_left: List[float] = field(default_factory=lambda: [
        -0.84, -2.03,  1.15,  1.15,  2.71,  1.60, -2.99, 888.00,
    ])
    init_state_right: List[float] = field(default_factory=lambda: [
         1.16,  2.01, -0.79, -0.68, -2.84, -1.61,  2.37, 832.00,
    ])

    # Realman SDK uses degrees for joint angles and meters for positions
    joint_units: List[str] = field(default_factory=lambda: [
        'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'degree', 'm',
    ])
    pose_units: List[str] = field(default_factory=lambda: [
        'm', 'm', 'm', 'degree', 'degree', 'degree', 'm',
    ])


@TeleoperatorConfig.register_subclass("bi_realman_leader_end_effector")
@dataclass
class BiRealmanLeaderEndEffectorConfig(BiRealmanLeaderConfig, BaseLeaderEndEffectorConfig):
    """
    Configuration for Bi-Realman leader with end effector.
    """

    pass