# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import numpy as np

from glob import glob

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
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, crawl

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

def get_depth(filename):
    depth = read_depth(filename, div=1000)
    return depth
    
def get_mask(filename):
    return read_image(filename, '1')
    
#######################################################

def get_sequences(args):
    seqs = crawl(args.src, 'metadata.json')
    seqs = [os.path.dirname(seq) for seq in seqs]
    return seqs


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

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

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

    ### Initialize lists and dicts
    videos = glob(f'{seq}/*.mp4')
    cameras = [os.path.basename(v).split('-')[0].split('_')[0] for v in videos]
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}

    ############ LOOP OVER CAMERAS
    for cam in cameras:

        from anydata.utils.write import write_jpg_from_mp4
        video = f'{seq}/{cam}_camera-images-rgb.mp4'
        # print('@@@@', f'{dst}/rgb/{cam}')
        # resolution, length = write_jpg_from_mp4(f'{dst}/rgb/{cam}', video)
        # print(resolution, length)

        left_robot_pose = f'{seq}/left_robot_pose.npy'
        right_robot_pose = f'{seq}/right_robot_pose.npy'

        print()
        print(robot_pose)
        robot_pose = np.load(robot_pose)
        for i in range(len(robot_pose)):
            print(i, robot_pose[i].shape)

        import sys
        sys.exit()

        ######## RGB FILENAMES
        for i, filename_rgb in enumerate(filename_rgbs):
            obj, num, _, _ = filename_rgb.split('/')[-4:]
            frame = frame_name(i)

            ######## RGB
            if 'rgb' not in labels:
                labels.append('rgb')
            rgb = np.array(read_image(filename_rgb))
            filename_rgb_out = f'{dst}/rgb/{cam}/{frame}.jpg'
            write_image(filename_rgb_out, rgb)

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

            ######## INTRINSICS + EXTRINSICS
            if 'intrinsics' not in labels:
                labels.append('intrinsics')
            if 'extrinsics' not in labels:
                labels.append('extrinsics')
            lowdim[filename_lowdim]['extrinsics'] = extrinsics[i]
            lowdim[filename_lowdim]['intrinsics'] = intrinsics

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

        ######## DEPTH FILENAMES
        for i, filename_depth in enumerate(filename_depths):
            obj, num, _, _ = filename_depth.split('/')[-4:]
            frame = frame_name(i)

            ######## DEPTH
            if 'depth' not in labels:
                labels.append('depth')
            depth = get_depth(filename_depth)
            filename_depth_out = f'{dst}/depth/{cam}/{frame}.npz'
            write_npz(filename_depth_out, dict(depth=depth))

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

        ######## MASK DEPTH FILENAMES
        for i, filename_mask_depth in enumerate(filename_mask_depths):
            obj, num, _, _ = filename_mask_depth.split('/')[-4:]
            frame = frame_name(i)

            ######## MASK DEPTH
            if 'mask_depth' not in labels:
                labels.append('mask_depth')
            mask_depth = get_mask(filename_mask_depth)
            filename_mask_depth_out = f'{dst}/mask_depth/{cam}/{frame}.png'
            write_image(filename_mask_depth_out, mask_depth)

        if len(filename_mask_depths) > 0:
            resolution[cam]['mask_depth'] = mask_depth.size[::-1]
            num_frames[cam]['mask_depth'] = len(filename_mask_depths)

    ######## 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='WildRGBD',
            tags=['real','static','inward','object'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=fps,
        rgb=dict(extension='jpg'),
        intrinsics=dict(model='pinhole'),
        extrinsics=dict(transform='cam2world',metric=False),
        depth=dict(extension='npz',metric=False,sparse=True),
        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)

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