import pathlib

### Task parameters
DATA_DIR = '/data0/hongyi/act_data'
SIM_TASK_CONFIGS = {
    'sim_open_drawer_real30_dream87_fh': {     # 30 real + 9 (b1) + 45 (b2) + 33 (b3: 15/10/8) OSCAR fixed-horizon dreams, stateless
        'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream87_fh',
        'num_episodes': 117,
        'episode_len': 200,           # longest demo is 192 (dreams); pad to 200
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_open_drawer_real30_dream54_fh': {     # 30 real + 9 (batch1) + 45 (batch2) OSCAR fixed-horizon dreams, stateless
        'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream54_fh',
        'num_episodes': 84,
        'episode_len': 200,           # longest demo is 192 (dreams); pad to 200
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_open_drawer_real30_dream9_fh': {      # 30 real + 9 OSCAR fixed-horizon dreams (3/room, rooms 2/4/5), stateless
        'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream9_fh',
        'num_episodes': 39,
        'episode_len': 200,           # longest demo is 192 (dreams); pad to 200
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_open_drawer_real30_dream8': {         # 30 real + 8 OSCAR dreams (room2), stateless
        'dataset_dir': '/root/pwm_local/data0/act_data_mix/open_drawer_real30_dream8',
        'num_episodes': 38,
        'episode_len': 200,           # longest demo is 177; pad to 200
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_open_drawer_100demos': {
        'dataset_dir': DATA_DIR + '/open_drawer_100demos',
        'num_episodes': 100,
        'episode_len': 200,           # longest demo is 177; pad to 200
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 50,             # ~25% of episode_len, matches ALOHA's 100/400 ratio
    },

    'sim_push_box_100demos': {
        'dataset_dir': DATA_DIR + '/push_box_100demos',
        'num_episodes': 100,
        'episode_len': 320,           # longest demo is 307; pad to 320
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,             # matches open_drawer_100demos_chunk25_30k recipe
    },

    'sim_push_box_50demos': {
        # First 50 of the 100 push-box demos (num_episodes=N loads episode_0..N-1).
        # Deliberately WEAKER BC base — to test whether WM RL has more headroom to
        # improve when the BC starting point is lower than the 100-demo (80%) policy.
        'dataset_dir': DATA_DIR + '/push_box_100demos',
        'num_episodes': 50,
        'episode_len': 320,
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_push_box_100demos_table2': {
        # Same demos (action/qpos identical) but RGB rendered at table_idx=2 — so the
        # ACT is native to the table the Push-Box WM was built at (robot_new/Push-Box_halffail).
        'dataset_dir': DATA_DIR + '/push_box_100demos_table2',
        'num_episodes': 100,
        'episode_len': 320,
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_open_laptop_100demos': {
        # Built directly from EgoVLA_SIM/Open-Laptop hdf5 (build_act_from_egovla.py).
        # 38-d upper-body proprio/action (matches act/eval_act_in_sim.py).
        'dataset_dir': DATA_DIR + '/open_laptop_100demos',
        'num_episodes': 100,
        'episode_len': 160,           # longest demo 146; pad to 160
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_pour_balls_100demos': {
        'dataset_dir': DATA_DIR + '/pour_balls_100demos',
        'num_episodes': 100,
        'episode_len': 250,           # longest demo 236; pad to 250
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_open_drawer_30demos': {
        # First 30 (sorted) EgoVLA_SIM/Open-Drawer demos via build_act_from_egovla.py.
        'dataset_dir': DATA_DIR + '/open_drawer_30demos',
        'num_episodes': 30,
        'episode_len': 200,           # longest demo 177; pad to 200
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_open_laptop_30demos': {
        'dataset_dir': DATA_DIR + '/open_laptop_30demos',
        'num_episodes': 30,
        'episode_len': 160,           # longest demo 146; pad to 160
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_pour_balls_30demos': {
        'dataset_dir': DATA_DIR + '/pour_balls_30demos',
        'num_episodes': 30,
        'episode_len': 200,           # longest demo 191; pad to 200
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_push_box_30demos': {
        'dataset_dir': DATA_DIR + '/push_box_30demos',
        'num_episodes': 30,
        'episode_len': 280,           # longest demo 267; pad to 280
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_multitask_30each': {         # 4 tasks x 30 demos: open_drawer, open_laptop, pour_balls, push_box (120 eps)
        'dataset_dir': DATA_DIR + '/multitask_30each',
        'num_episodes': 120,
        'episode_len': 280,           # max len over the 4 (push_box 267); pad to 280
        'camera_names': ['top'],
        'state_dim': 38,
        'action_dim': 38,
        'chunk_size': 25,
    },

    'sim_transfer_cube_scripted':{
        'dataset_dir': DATA_DIR + '/sim_transfer_cube_scripted',
        'num_episodes': 50,
        'episode_len': 400,
        'camera_names': ['top']
    },

    'sim_transfer_cube_human':{
        'dataset_dir': DATA_DIR + '/sim_transfer_cube_human',
        'num_episodes': 50,
        'episode_len': 400,
        'camera_names': ['top']
    },

    'sim_insertion_scripted': {
        'dataset_dir': DATA_DIR + '/sim_insertion_scripted',
        'num_episodes': 50,
        'episode_len': 400,
        'camera_names': ['top']
    },

    'sim_insertion_human': {
        'dataset_dir': DATA_DIR + '/sim_insertion_human',
        'num_episodes': 50,
        'episode_len': 500,
        'camera_names': ['top']
    },
}

### Simulation envs fixed constants
DT = 0.02
JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"]
START_ARM_POSE = [0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239,  0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239]

XML_DIR = str(pathlib.Path(__file__).parent.resolve()) + '/assets/' # note: absolute path

# Left finger position limits (qpos[7]), right_finger = -1 * left_finger
MASTER_GRIPPER_POSITION_OPEN = 0.02417
MASTER_GRIPPER_POSITION_CLOSE = 0.01244
PUPPET_GRIPPER_POSITION_OPEN = 0.05800
PUPPET_GRIPPER_POSITION_CLOSE = 0.01844

# Gripper joint limits (qpos[6])
MASTER_GRIPPER_JOINT_OPEN = 0.3083
MASTER_GRIPPER_JOINT_CLOSE = -0.6842
PUPPET_GRIPPER_JOINT_OPEN = 1.4910
PUPPET_GRIPPER_JOINT_CLOSE = -0.6213

############################ Helper functions ############################

MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_POSITION_CLOSE) / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_POSITION_CLOSE) / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE
PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE
MASTER2PUPPET_POSITION_FN = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(MASTER_GRIPPER_POSITION_NORMALIZE_FN(x))

MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)
PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)
MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x))

MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)

MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE))
PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE))

MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2