# Copyright 2026 Toyota Research Institute.  All rights reserved.
# Inspired by: https://github.com/physical-superintelligence-lab/Humanoid-Everyday

import os
import cv2
import lzma
import numpy as np
import h5py

from PIL import Image
from glob import glob
from tqdm import tqdm

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

from externals.egodex.simple_dataset import SimpleDataset

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

def get_rgb(filename):
    with open(filename, "rb") as f:
        img = Image.open(f)
        img = img.convert("RGB")
        return np.array(img)

def get_depth(filename):
    with open(filename, "rb") as f:
        compressed_data = f.read()
        decompressed = lzma.decompress(compressed_data)
        depth_array = np.frombuffer(decompressed, dtype=np.uint16).reshape((480, 640)) / 1000
        return depth_array
    
#######################################################

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


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

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

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

    dataset = SimpleDataset(seq)
    episodes = [int(l) for l in dataset.cumulative_len]

    dsts, cnt = [], 0
    for ep in tqdm(range(len(episodes)), ncols=96, leave=False):
        dst = f'{dst_all}/%04d' % ep

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

        cam = cameras[0]
        dense = {label: dict() for label in dense_labels}

        while cnt < dataset.cumulative_len[ep]:

            epid, frid, tfs, cam_ext, cam_int, rgb, lang_instruct, confs = dataset[cnt]
            frame = frame_name(frid)
            cnt += 1

            ######## RGB
            rgb = rgb.permute(1, 2, 0).numpy()
            dense['rgb'][frame] = rgb

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

            ######## INTRINSICS
            lowdim[filename_lowdim]['intrinsics'] = cam_int
            lowdim[filename_lowdim]['extrinsics'] = cam_ext

            ######## ACTION
            lowdim[filename_lowdim]['action'] = {
                'action': tfs, 
                'confidence': confs,
            }

        ######## LANGUAGE
        language = dict(prompt=[lang_instruct])

        ######## 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='EgoDex',
                tags=['real','dynamic','egocentric','human'],
                raw_id=seq.replace(f'{args.src}/', ''),
            ),
            labels=labels,
            cameras=cameras,
            resolution=resolution,
            num_frames=num_frames,
            framerate=30,
            rgb=dict(extension='jpg'),
            intrinsics=dict(model='pinhole'),
            extrinsics=dict(transform='cam2world',metric=True),
            depth=None,
            semantic=None,
            action=dict(format='EgoDex'),
            language=language,
            specific=None,
        )
        write_json(filename, seq_metadata)
        dsts.append(dst)

    return dsts

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

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

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