# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import cv2
import numpy as np

import torch
from glob import glob

from anydata.utils.geometry import invert_extrinsics
from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth, read_json
from anydata.utils.write import write_image, write_npz, write_json, write_png8
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name
from scipy.spatial.transform import Rotation as R

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

def vect2trans(vec):
    vec = torch.tensor(vec)
    if len(vec.shape) == 1:
        vec = vec.unsqueeze(0)
    rot = torch.tensor(R.from_rotvec(vec[:, :3]).as_matrix())
    trans = vec[:, 3:].unsqueeze(-1)
    T = torch.cat([rot, trans], -1)
    extrinsics = torch.cat([T, torch.tensor([[[0.0, 0.0, 0.0, 1.0]]])], 1).float()
    return extrinsics[0].numpy()

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

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


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args)

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

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

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

##################
    dirname = os.path.dirname(seq)
    name = os.path.basename(seq).split('_')[1]
    files = sorted(glob(f'{dirname}/log_{name}_*.npz'))

    for i, file in enumerate(files):
        for cam in cameras:
            frame = frame_name(i)

            data = np.load(file)
            keys = list(data.keys())

            extrinsics = data['cam_pose']
            extrinsics = vect2trans(extrinsics)

            global_extrinsics = data['global_T_local_pose']
            global_extrinsics = vect2trans(global_extrinsics)

            extrinsics = global_extrinsics @ extrinsics

            fx, fy, cx, cy = data['cam_intrinsics']
            intrinsics = np.array([
                [ fx, 0.0,  cx],
                [0.0,  fy,  cy],
                [0.0, 0.0, 1.0]
            ])

            ######## RGB
            if 'rgb' not in labels: labels.append('rgb')
            rgb = data['rgb']
            filename_rgb_out = f'{dst}/rgb/{cam}/{frame}.jpg'
            write_image(filename_rgb_out, rgb)

            ######## DEPTH
            if 'depth' not in labels: labels.append('depth')
            depth = data['depth'] / data['depth_scale']
            filename_depth_out = f'{dst}/depth/{cam}/{frame}.png'
            write_png8(filename_depth_out, depth)

            ######## LOWDIM RGB             
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['camera'] = cam
            lowdim[filename_lowdim]['timestep'] = data['timestamp']

            ######## INTRINSICS + EXTRINSICS
            if 'intrinsics' not in labels: labels.append('intrinsics')
            if 'extrinsics' not in labels: labels.append('extrinsics')
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['extrinsics'] = extrinsics
            lowdim[filename_lowdim]['intrinsics'] = intrinsics

            ######## ACTION
            if 'action' not in labels: labels.append('action')
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            action_keys = {key: val for key, val in data.items() if key. endswith('action')}
            pose_keys = {key: val for key, val in data.items() if key. endswith('pose')}
            tactile_keys = {key: val for key, val in data.items() if key. endswith('wrench')}
            lowdim[filename_lowdim]['action'] = {**action_keys, **pose_keys, **tactile_keys}

    if len(files) > 0:
        resolution[cam]['rgb'] = rgb.shape[:2]
        num_frames[cam]['rgb'] = len(files)
        resolution[cam]['depth'] = depth.shape[:2]
        num_frames[cam]['depth'] = len(files)

    ######## 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='TZK',
            tags=['real','robotics','indoor'],
            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=dict(extension='png8',metric=True,sparse=True),
        semantic=None,
        action=dict(format='TZK'),
        language=dict(prompt=["You are a helpful robot assistant finishing tasks to help people's daily lives."]),
        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)

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