# 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

from PIL import Image
from glob import glob

from anydata.utils.read import read_json
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_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, 'color')
    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
    cameras = ['0']
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}
    dense_labels = ['rgb','depth']

    ### Get robot data
    data = read_json(f'{seq}/data.json')
    robot_type = data[0]['robot_type'].upper() if 'robot_type' in data[0] else 'H1'
    has_depth = False

    ### Intrinsics
    if robot_type == 'G1':
        fx, fy = 389.07278, 389.07278
        cx, cy = 321.61887, 238.43630
    elif robot_type == 'H1':
        fx, fy = 392.03189, 392.03189
        cx, cy = 320.19580, 235.58174
    else:
        raise ValueError(f'Invalid robot {robot_type}')
    intrinsics = np.array([
        [fx, 0., cx],
        [0., fy, cy], 
        [0., 0., 1.],
    ])

    ### Get language prompts 
    language = None
    tasks = seq.split('/')[-4:-1]
    for task in tasks:
        task = task.replace(' ', '_')
        if task in args.language.keys():
            language = dict(
                task=task.replace('_', ' '),
                prompt=[args.language[task]],
                category=seq.split('/')[4],
            )
            break
    assert language is not None, 'Task not found in language dictionary'

    ############ LOOP OVER CAMERAS
    for cam in cameras:
        dense = {label: dict() for label in dense_labels}

        ######## ACTION DATA
        for i in range(len(data)):
            frame = frame_name(i)

            filename_rgb = seq + '/' + data[i]['image']
            filename_depth = seq + '/' + data[i]['depth']

            ######## RGB
            rgb = get_rgb(filename_rgb)
            dense['rgb'][frame] = rgb

            ######## LOWDIM RGB
            prepare_lowdim(lowdim, dst, cam, frame)

            ######## INTRINSICS
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['intrinsics'] = intrinsics

            ######## DEPTH
            if os.path.exists(filename_depth):
                depth = get_depth(filename_depth)
                dense['depth'][frame] = depth
                has_depth = True

            ######## ACTION
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['action'] = {
                'actions': data[i]['actions'],
                'states': data[i]['states'],
            }

        ######## 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='HumanoidEveryday',
            tags=['real','dynamic','egocentric','robotics','humanoid'],
            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=None,
        depth=dict(extension='npz',metric=True,sparse=True) if has_depth else None,
        semantic=None,
        action=dict(format='HumanoidEveryday'),
        language=language,
        specific=dict(
            robot_type=robot_type,
        ),
    )
    write_json(filename, seq_metadata)

    return dst

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

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

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