# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import json
import numpy as np

from glob import glob

from anydata.utils.write import write_image, write_npz, write_json
from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth, read_json
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name
from anydata.utils.geometry import invert_extrinsics
from anydata.utils.colmap import qvec2rotmat

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

def get_sequences(args):
    seqs = glob(f'{args.src}/**/*.mcap', recursive=True)
    return seqs


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

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

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

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

    ### Initializer MCAP fields
    images = reader.get_image_topics()
    all_intrinsics = reader.get_all_intrinsics()
    all_extrinsics = reader.get_all_extrinsics()
    all_lowdim = reader.get_all_lowdim()

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

    # Map cameras to extrinsics
    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',
    }

    # Which keys are actions
    action_keys = [
        '/joint_cmd_hand_right__right_hand_index_0_joint',
    ]

    # Loop over cameras
    for c, image in enumerate(images):
        # Get images
        timestamps, rgbs = reader.get_images(image)
        # Loop over all images
        for i in range(len(rgbs)):
            frame = frame_name(i)
            cam = cameras[c] # Camera name

            # Save image
            filename_rgb_out = f'{dst}/rgb/{cam}/{frame}.jpg'
            write_image(filename_rgb_out, rgbs[i])

            # Get intrinsics
            info = all_intrinsics[f'/{cam}/color/camera_info']
            distortion = info['distortion_model']
            K, D = info['K'], info['D']
            intrinsics = K

            # Get extrinsics
            extrinsics = all_extrinsics[extrinsics_keys[cam]]['transform_matrix']

            # Save intrinsics and extrinsics
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['intrinsics'] = intrinsics
            lowdim[filename_lowdim]['extrinsics'] = extrinsics

            # Save action fields
            lowdim[filename_lowdim]['action'] = {
                action_keys[0]: all_lowdim[action_keys[0]][0]
            }

    # Log resolution and number of frames
    resolution[cam]['rgb'] = rgbs[0].shape[:2]
    num_frames[cam]['rgb'] = len(rgbs)

    ######## WRITE LOWDIM
    for key, val in lowdim.items():
        write_npz(key, val)

############ METADATA 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='MCAP',
            tags=['sim','dynamic','driving'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=10,
        rgb=dict(extension='jpg'),
        intrinsics=dict(model='pinhole'),
        extrinsics=dict(transform='cam2world',metric=True),
        depth=None,
        semantic=None,
        action=None,
        language=None,
        specific=None,
    )
    write_json(filename, seq_metadata)

    return dst

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

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

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