# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import torch
import numpy as np

from glob import glob

from anydata.utils.geometry import pose_to_matrix
from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth
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

import tensorflow_datasets as tfds

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

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=str, nargs='+')
    parser.add_argument("--num_procs", type=int, default=16)
    parser.add_argument("--local_folder", type=str, default='/data/cv_unified')
    args = parser.parse_args()

    args.src, args.dst = args.path
    args.dst = f'{args.local_folder}/{args.dst}'

    return args

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

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 = glob(f'{args.src}/*')
    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 = ['head','wrist_left','wrist_right']

    dataset_name, data_dir = os.path.basename(seq), os.path.dirname(seq)
    ds = tfds.load(dataset_name, split='train', data_dir=data_dir)
    for i, episode in enumerate(ds.take(len(ds))):

        num_frames = {cam: dict() for cam in cameras}
        resolution = {cam: dict() for cam in cameras}
        labels, lowdim = [], {}

        metadata = str(episode['episode_metadata']['file_path'])
        metadata = metadata.split("'")[1]
        name = '/'.join(metadata.split('/')[4:])
        dst_name = f'{dst}/{name}'

        for j, step in enumerate(episode['steps']):
            frame = frame_name(j)

            for cam in cameras:

                ######## RGB
                if 'rgb' not in labels: labels.append('rgb')
                rgb = step['observation'][f'image_camera_{cam}'].numpy()
                filename_rgb = f'{dst_name}/rgb/{cam}/{frame}.jpg'
                write_image(filename_rgb, rgb)

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

                resolution[cam]['rgb'] = rgb.shape[:2]
                num_frames[cam]['rgb'] = j + 1

                ######## DEPTH
                if cam.startswith('wrist'):
                    if 'depth' not in labels: labels.append('depth')
                    depth = step['observation'][f'depth_camera_{cam}'].numpy().astype(np.float32) / 1000
                    filename_depth = f'{dst_name}/depth/{cam}/{frame}.png'
                    write_png8(filename_depth, depth)

                    resolution[cam]['depth'] = depth.shape[:2]
                    num_frames[cam]['depth'] = j + 1

                ######## ACTION             
                filename_lowdim = add_key_to_dict(lowdim, f'{dst_name}/lowdim/{cam}/{frame}.npz')
                lowdim[filename_lowdim]['action'] = {
                    'action': step['action'], 
                    'state': {key: val for key, val in step['observation'].items() \
                        if not key.startswith('image') and not key.startswith('depth')
                    },
                }

        ######## LANGUAGE
        if 'language' not in labels: labels.append('language')
        instruction = step['language_instruction'].numpy().decode('utf-8')
        task, _, prompt = instruction.split('@')

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

    ############ METADATA 
        filename = f'{dst_name}/metadata.json'
        seq_metadata = fill_metadata(
            args=args,
            info=dict(
                name='Galaxea',
                tags=['real','dynamic','robotics','tabletop'],
                raw_id=seq.replace(f'{args.src}/', '') + f'/{name}',
            ),
            labels=labels,
            cameras=cameras,
            resolution=resolution,
            num_frames=num_frames,
            framerate=10,
            rgb=dict(extension='jpg'),
            intrinsics=None,
            extrinsics=None,
            depth=dict(extension='png8'),
            semantic=None,
            action=dict(format='Galaxea'),
            language=dict(task=task, prompt=[prompt]),
            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)

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