# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import h5py
import numpy as np

from glob import glob

from anydata.utils.read import read_numpy, 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

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

def get_sequences(args):
    seqs = crawl(args.src, '*.hdf5')
    return seqs


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

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

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

    with h5py.File(seq, "r") as data:

        ### Initialize lists and dicts
        cameras = [d for d in list(data['obs'].keys()) if d.endswith('_image')]
        num_frames = {cam: dict() for cam in cameras}
        resolution = {cam: dict() for cam in cameras}
        labels, lowdim = [], {}
        dense_labels = ['rgb']

        dsts, cnt = [], 0
        while cnt < data['obs'][cameras[0]].shape[0]:

            dst = f'{dst_all}/%02d' % len(dsts)
            dsts.append(dst)

            for j, cam in enumerate(cameras):
                dense = {label: dict() for label in dense_labels}

                rgbs = data['obs'][cam]

                ######## RGB FILENAMES
                for i in range(cnt, rgbs.shape[0]):
                    frame = frame_name(i)
                    cnt += 1

                    ######## RGB
                    rgb = np.transpose(rgbs[i, 0], (1, 2, 0))
                    dense['rgb'][frame] = (rgb * 255).astype(np.uint8)

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

                    ######## ACTION
                    lowdim[filename_lowdim]['action'] = {
                        **{key: data[key][i, 0] for key in ['action','reward','task_done']},
                        **{key: val[i, 0] for key, val in data['obs'].items() if key.startswith('robot0')},
                    }

                    # print(lowdim[filename_lowdim]['action'].keys())
                    
                    if data['task_done'][i, 0]:
                        break

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

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

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

    return dsts

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

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

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