# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import numba
import numpy as np

from glob import glob
from scipy.spatial.transform import Rotation

from anydata.utils.read import read_json, read_yaml
from anydata.utils.write import write_json, write_lowdim, write_labels
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, crawl, prepare_lowdim

#######################################################

# D435 depth-to-color extrinsic (factory calibration, consistent across units).
# The URDF d435_link is at the depth sensor origin; the color camera is offset
# ~15mm along the sensor X axis. Verified on two D435 units: 14.9mm (HE) and
# 14.7mm (MCAP). Rotation is < 1 degree, approximated as identity.
D435_DEPTH_TO_COLOR = np.eye(4)
D435_DEPTH_TO_COLOR[:3, 3] = [0.0149, 0.0, 0.0]

EXTRINSICS_KEYS = {
    'head_camera': 'd435_link->head_camera_link',
    'left_wrist_camera': 'left_wrist_yaw_link->left_wrist_camera_mount_link',
    'right_wrist_camera': 'right_wrist_yaw_link->right_wrist_camera_mount_link',
}

# ROS optical frame convention: maps optical coords -> link coords
T_LINK_FROM_OPTICAL = np.array(
    [[0, 0, 1, 0], [-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=float
)
T_OPTICAL_FROM_LINK = np.linalg.inv(T_LINK_FROM_OPTICAL)

# URDF chain pelvis -> d435_link (at zero config, static transforms only).
# These are chained in process_sequence from MCAP extrinsics.
URDF_CHAIN_PELVIS_TO_D435 = [
    'pelvis->waist_yaw_link',
    'waist_yaw_link->waist_roll_link',
    'waist_roll_link->torso_link',
    'torso_link->d435_link',
]

# D435 depth-to-color with full rotation (factory calibration, from RealSense)
D435_DEPTH_TO_COLOR_FULL = np.eye(4)
D435_DEPTH_TO_COLOR_FULL[:3, :3] = np.array([
    [0.999901533126831, -0.012537548318505287, 0.006302413530647755],
    [0.012552149593830109, 0.9999186396598816, -0.002282573375850916],
    [-0.006273282691836357, 0.002361457562074065, 0.9999775290489197],
])
D435_DEPTH_TO_COLOR_FULL[:3, 3] = [0.01490413211286068, -5.230475835560355e-06, 0.00011887117580045015]


# Torque command joints (no hand joints)
TORQUE_JOINTS = [
    'left_shoulder_pitch_joint', 'left_shoulder_roll_joint', 'left_shoulder_yaw_joint',
    'left_elbow_joint', 'left_wrist_roll_joint', 'left_wrist_pitch_joint', 'left_wrist_yaw_joint',
    'right_shoulder_pitch_joint', 'right_shoulder_roll_joint', 'right_shoulder_yaw_joint',
    'right_elbow_joint', 'right_wrist_roll_joint', 'right_wrist_pitch_joint', 'right_wrist_yaw_joint',
    'waist_yaw_joint', 'waist_roll_joint', 'waist_pitch_joint',
    'left_hip_pitch_joint', 'left_hip_roll_joint', 'left_hip_yaw_joint',
    'left_knee_joint', 'left_ankle_pitch_joint', 'left_ankle_roll_joint',
    'right_hip_pitch_joint', 'right_hip_roll_joint', 'right_hip_yaw_joint',
    'right_knee_joint', 'right_ankle_pitch_joint', 'right_ankle_roll_joint',
]

# Joint state topics (current positions)
JOINT_STATE_JOINTS = [
    'left_shoulder_pitch_joint', 'left_shoulder_roll_joint', 'left_shoulder_yaw_joint',
    'left_elbow_joint', 'left_wrist_roll_joint', 'left_wrist_pitch_joint', 'left_wrist_yaw_joint',
    'right_shoulder_pitch_joint', 'right_shoulder_roll_joint', 'right_shoulder_yaw_joint',
    'right_elbow_joint', 'right_wrist_roll_joint', 'right_wrist_pitch_joint', 'right_wrist_yaw_joint',
    'waist_yaw_joint', 'waist_roll_joint', 'waist_pitch_joint',
    'left_hip_pitch_joint', 'left_hip_roll_joint', 'left_hip_yaw_joint',
    'left_knee_joint', 'left_ankle_pitch_joint', 'left_ankle_roll_joint',
    'right_hip_pitch_joint', 'right_hip_roll_joint', 'right_hip_yaw_joint',
    'right_knee_joint', 'right_ankle_pitch_joint', 'right_ankle_roll_joint',
    'left_hand_index_0_joint', 'left_hand_index_1_joint',
    'left_hand_middle_0_joint', 'left_hand_middle_1_joint',
    'left_hand_thumb_0_joint', 'left_hand_thumb_1_joint', 'left_hand_thumb_2_joint',
    'right_hand_index_0_joint', 'right_hand_index_1_joint',
    'right_hand_middle_0_joint', 'right_hand_middle_1_joint',
    'right_hand_thumb_0_joint', 'right_hand_thumb_1_joint', 'right_hand_thumb_2_joint',
]

#######################################################

# def get_depth(topic):
#     # if topic == '/head_camera/zed_node/left/image_rect_color/compressed':
#     #     key = '/head_camera/zed_node/left/camera_info'
#     # elif topic == '/head_camera/zed_node/right/image_rect_color/compressed':
#     #     key = '/head_camera/zed_node/right/camera_info'
#     if topic == '/left_wrist_camera/color/image_rect_raw/compressed':
#         key = '/left_wrist_camera/depth/image_rect_raw'
#     elif topic == '/right_wrist_camera/color/image_rect_raw/compressed':
#         key = '/right_wrist_camera/depth/image_rect_raw'
#     else:
#         raise ValueError('Invalid topic')

# RGB image topic → camera_info topic. Covers both the older D435 head camera
# (used by early teleop episodes) and the newer ZED stereo head camera (used by
# rollout / move_block_on_plate / later episodes), plus the two D435 wrist
# cameras that are constant across both eras.
INTRINSICS_TOPIC_FOR_RGB = {
    '/head_camera/color/image_raw/compressed':                  '/head_camera/color/camera_info',
    '/head_camera/zed_node/left/image_rect_color/compressed':   '/head_camera/zed_node/left/camera_info',
    '/head_camera/zed_node/right/image_rect_color/compressed':  '/head_camera/zed_node/right/camera_info',
    '/left_wrist_camera/color/image_rect_raw/compressed':       '/left_wrist_camera/color/camera_info',
    '/right_wrist_camera/color/image_rect_raw/compressed':      '/right_wrist_camera/color/camera_info',
}

# RGB image topic → static-transform key in the MCAP tf broadcast.
# All entries are sensor-link → color-optical-frame, i.e. the small rotation
# that puts the image plane in the OpenCV convention (X right, Y down, Z fwd).
EXTRINSICS_KEY_FOR_RGB = {
    '/head_camera/color/image_raw/compressed':                  'd435_link->head_camera_link',
    '/head_camera/zed_node/left/image_rect_color/compressed':   'head_camera_left_camera_frame->head_camera_left_camera_optical_frame',
    '/head_camera/zed_node/right/image_rect_color/compressed':  'head_camera_right_camera_frame->head_camera_right_camera_optical_frame',
    '/left_wrist_camera/color/image_rect_raw/compressed':       'left_wrist_camera_link->left_wrist_camera_color_optical_frame',
    '/right_wrist_camera/color/image_rect_raw/compressed':      'right_wrist_camera_link->right_wrist_camera_color_optical_frame',
}


def get_intrinsics(topic, all_intrinsics):
    key = INTRINSICS_TOPIC_FOR_RGB.get(topic)
    if key is None:
        raise ValueError(f'Unknown RGB topic for intrinsics dispatch: {topic}')
    data = all_intrinsics.get(key)
    return data['K'] if data is not None else None


def get_extrinsics(topic, all_extrinsics):
    key = EXTRINSICS_KEY_FOR_RGB.get(topic)
    if key is None:
        raise ValueError(f'Unknown RGB topic for extrinsics dispatch: {topic}')
    data = all_extrinsics.get(key)
    return data['transform_matrix'] if data is not None else None

def get_states_and_actions(all_lowdim, t):

    # Get all joint positions and commands as full vectors
    joint_states_all = get_joint_state_vector(
        all_lowdim, '/joint_states', JOINT_STATE_JOINTS, t)
    joint_cmd_all = get_joint_state_vector(
        all_lowdim, '/joint_cmd_all', JOINT_STATE_JOINTS, t)

    # Reorder from URDF [index0, index1, middle0, middle1, thumb0, thumb1, thumb2]
    # to HumanoidEveryday/Unitree SDK [thumb0, thumb1, thumb2, middle0, middle1, index0, index1]
    _MCAP_TO_HE_HAND = [4, 5, 6, 2, 3, 0, 1]

    # --- states (measured feedback, matches HumanoidEveryday format) ---
    states = {}

    # arm_state: left arm (7) + right arm (7) = 14
    states['arm_state'] = np.concatenate([
        joint_states_all[0:7], joint_states_all[7:14]
    ])
    # leg_state: waist (3) + left leg (6) + right leg (6) = 15
    states['leg_state'] = np.concatenate([
        joint_states_all[14:17], joint_states_all[17:23], joint_states_all[23:29]
    ])
    # hand_state: left hand (7) + right hand (7) = 14
    # Reorder from URDF to HumanoidEveryday/Unitree SDK order
    states['hand_state'] = np.concatenate([
        joint_states_all[29:36][_MCAP_TO_HE_HAND],
        joint_states_all[36:43][_MCAP_TO_HE_HAND],
    ])

    # imu: dict with quaternion, accelerometer, gyroscope, rpy
    imu_raw = get_lowdim_at_time(all_lowdim, '/imu', t)
    if imu_raw is not None:
        imu_flat = imu_raw.flatten().astype(np.float64)
        quat_xyzw = imu_flat[0:4]
        rpy = Rotation.from_quat(quat_xyzw).as_euler('xyz')
        states['imu'] = {
            'quaternion': quat_xyzw,
            'accelerometer': imu_flat[7:10],
            'gyroscope': imu_flat[4:7],
            'rpy': rpy,
        }

    # --- actions (commands/targets, matches HumanoidEveryday format) ---
    actions = {}

    # left_angles / right_angles: hand/finger joint commands (7 each)
    actions['left_angles'] = joint_cmd_all[29:36][_MCAP_TO_HE_HAND]
    actions['right_angles'] = joint_cmd_all[36:43][_MCAP_TO_HE_HAND]

    # sol_q: arm joint commands (14: 7 left + 7 right)
    actions['sol_q'] = np.concatenate([
        joint_cmd_all[0:7], joint_cmd_all[7:14]
    ])

    # tau_ff: torque commands (29 whole-body; first 14 are arm torques)
    actions['tau_ff'] = get_joint_state_vector(
        all_lowdim, '/joint_cmd_whole_body_torque', TORQUE_JOINTS, t)

    # left_pose / right_pose: EE poses as 4x4 matrices
    left_pose = get_ee_pose_matrix(all_lowdim, '/current_left_hand_ee_link', t)
    if left_pose is not None:
        actions['left_pose'] = left_pose
    right_pose = get_ee_pose_matrix(all_lowdim, '/current_right_hand_ee_link', t)
    if right_pose is not None:
        actions['right_pose'] = right_pose

    # EE target poses (G1-specific, no HumanoidEveryday equivalent)
    left_target = get_ee_pose_matrix(all_lowdim, '/ee_target_left', t)
    if left_target is not None:
        actions['left_target_pose'] = left_target
    right_target = get_ee_pose_matrix(all_lowdim, '/ee_target_right', t)
    if right_target is not None:
        actions['right_target_pose'] = right_target

    # leg commands (G1-specific, balance controller output)
    actions['leg_cmd'] = np.concatenate([
        joint_cmd_all[14:17], joint_cmd_all[17:23], joint_cmd_all[23:29]
    ])

    return states, actions

#######################################################

def get_lowdim_at_time(all_lowdim, topic_key, target_time):
    """Get nearest value from a lowdim topic at a given timestamp."""
    if topic_key not in all_lowdim:
        return None
    timestamps, values = all_lowdim[topic_key]
    idx = np.argmin(np.abs(timestamps - target_time))
    return values[idx]

def get_joint_state_vector(all_lowdim, prefix, joints, target_time):
    """Collect joint values into a single vector at a given timestamp."""
    values = []
    for joint in joints:
        key = f'{prefix}__{joint}'
        val = get_lowdim_at_time(all_lowdim, key, target_time)
        if val is not None:
            values.append(float(val.flatten()[0]))
        else:
            values.append(0.0)
    return np.array(values, dtype=np.float32)

def get_ee_pose_matrix(all_lowdim, prefix, target_time):
    """Get end-effector pose as a 4x4 transformation matrix at a given timestamp."""
    xyz = get_lowdim_at_time(all_lowdim, f'{prefix}__xyz', target_time)
    quat = get_lowdim_at_time(all_lowdim, f'{prefix}__quat', target_time)
    if xyz is None or quat is None:
        return None
    xyz = xyz.flatten()
    quat = quat.flatten()  # (x, y, z, w)
    rot = Rotation.from_quat(quat).as_matrix()
    T = np.eye(4, dtype=np.float64)
    T[:3, :3] = rot
    T[:3, 3] = xyz
    return T

#######################################################

def get_sequences(args):
    seqs = glob(f'{args.src}/**/*.mcap', recursive=True)
    seqs = [s for s in seqs if '/real/teleop/' in s]  # Only real teleop episodes
    return sorted(seqs)


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args, remext=True)

#######################################################

def process_sequence(i, seq, dst, args):

    ### Read info.yaml for metadata
    info_yaml_path = os.path.join(os.path.dirname(seq), 'info.yaml')
    info_yaml = read_yaml(info_yaml_path) if os.path.exists(info_yaml_path) else {}

    ### Get MCAP reader
    from anydata.converters.mcap_utils import MCAPReader
    reader = MCAPReader(seq)

    rgb_topics = reader.get_image_topics()
    depth_topics = reader.get_depth_topics()

    cameras = []
    for cam in rgb_topics:
        if 'head' in cam:    cam = '_'.join(cam.split('/')[1:4])
        elif 'wrist' in cam: cam = cam.split('/')[1]
        else: raise ValueError('Invalid camera name')
        cameras.append(cam)
    cam_map = {key: val for key, val in zip(rgb_topics, cameras)}

    ### Get all data
    all_intrinsics = reader.get_all_intrinsics()
    all_extrinsics = reader.get_all_extrinsics()
    all_lowdim = reader.get_all_lowdim()

    ### Initialize lists and dicts
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}
    dense_labels = ['rgb','depth']

    ### Language
    skill = info_yaml.get('skill', os.path.basename(os.path.dirname(os.path.dirname(os.path.dirname(seq)))))
    language = dict(
        task=skill,
        prompt=[skill.replace('_', ' ')],
    )

    ### Use head camera as reference timeline
    ref_cam = cameras[0]
    ref_topic = rgb_topics[0]
    ref_timestamps, ref_rgbs = reader.get_images(ref_topic)
    n_frames_total = len(ref_timestamps)

    ### Collect all camera images
    camera_data = {ref_topic: (ref_timestamps, ref_rgbs)}
    for cam in rgb_topics:
        if cam == ref_topic: continue
        if cam not in reader.get_image_topics(): continue
        timestamps, rgbs = reader.get_images(cam)
        # Sync to head camera timeline via nearest-neighbor
        indices = np.array([np.argmin(np.abs(timestamps - t)) for t in ref_timestamps])
        camera_data[cam] = (ref_timestamps, [rgbs[idx] for idx in indices])

    ### Loop over cameras
    for cam, (timestamps, rgbs) in camera_data.items():
        dense = {label: dict() for label in dense_labels}

        split = cam.split('/')
        mapped_cam = cam_map[cam]

        K = get_intrinsics(cam, all_intrinsics)
        T = get_extrinsics(cam, all_extrinsics)

        for frame_idx in range(n_frames_total):
            frame = frame_name(frame_idx)
            t = ref_timestamps[frame_idx]

            ######## RGB
            dense['rgb'][frame] = rgbs[frame_idx]

            ######## LOWDIM RGB
            prepare_lowdim(lowdim, dst, mapped_cam, frame)
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{mapped_cam}/{frame}.npz')

            ######## INTRINSICS
            if K is not None:
                lowdim[filename_lowdim]['intrinsics'] = K

            ######## EXTRINSICS
            if T is not None:
                lowdim[filename_lowdim]['extrinsics'] = T

            ######## STATES + ACTIONS
            states, actions = get_states_and_actions(all_lowdim, t)
            lowdim[filename_lowdim]['action'] = {'actions': actions, 'states': states}

        ######## WRITE LABELS
        write_labels(dst, mapped_cam, args.storage, dense, labels, resolution, num_frames)

    ######## WRITE LOWDIM
    write_lowdim(args, dst, labels, num_frames, lowdim)

    ######## METADATA
    robot = info_yaml.get('robot', 'g1')
    domain = info_yaml.get('domain', 'real')
    tags = ['dynamic', 'egocentric', 'humanoid', 'real', 'robotics']

    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='G1_MCAP',
            tags=tags,
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=30,
        rgb=dict(extension='jpg'),
        intrinsics=dict(model='pinhole'),
        extrinsics=dict(metric=True),
        depth=dict(extension='npz', metric=True, sparse=True),
        semantic=None,
        action=dict(format='G1_MCAP'),
        language=language,
        specific=dict(
            robot=robot,
            domain=domain,
            operator=info_yaml.get('operator'),
            station=info_yaml.get('station_name'),
        ),
    )
    write_json(filename, seq_metadata)

    return dst

#######################################################

if __name__ == '__main__':
    converter = os.path.basename(__file__)
    run(converter, get_sequences, parse_sequence, process_sequence)

#######################################################