# 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
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 = glob(f'{args.src}/**/observations.npz', recursive=True)
    seqs = [os.path.dirname(seq) for seq in seqs]
    return seqs


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args, remove=[0,1,2,6,9,11])

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

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

    ### Read data files
    observations_all = read_numpy(f'{seq}/observations.npz')
    extrinsics_all = read_numpy(f'{seq}/extrinsics.npz') if os.path.exists(f'{seq}/extrinsics.npz') else None
    intrinsics_all = read_numpy(f'{seq}/intrinsics.npz') if os.path.exists(f'{seq}/intrinsics.npz') else None
    action = read_numpy(f'{seq}/actions.npz')['actions']
    metadata = read_yaml(f'{seq}/metadata.yaml')

    ### Get observation information
    assert 'camera_id_to_semantic_name' in metadata.keys(), 'Camera IDs not found in metadata'
    tags = {v: k for k, v in metadata['camera_id_to_semantic_name'].items()}
    observations_keys = list(observations_all.keys())

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

    ### Filename template
    filename = "%s/%s/%s/%010d.%s"

    ### Get task from filename
    split = dst.split('/')
    task, station, sim_real, rollout_teleop, date = split[-6:-1]
    year, month, day = date.split('T')[0].split('_')
    depth_sparse = sim_real == 'real'

    ### Get language prompts 
    assert task in args.language.keys(), f'Task {task} not found in language dictionary'
    all_prompt = args.language[task]
    prompt = []
    for key in ['original','randomized','alternative']:
        if key in all_prompt:
            prompt.extend(all_prompt[key])

    ############ LOOP OVER CAMERAS
    for cam, val in tags.items():
        dense = {label: dict() for label in dense_labels}

        ######## RGB
        if val in observations_all.keys():
            rgbs = observations_all[val]

            for i in range(rgbs.shape[0]):
                frame = frame_name(i)
                dense['rgb'][frame] = rgbs[i]
                prepare_lowdim(lowdim, dst, cam, frame)

        ######## DEPTH
        val_depth = f'{val}_depth'
        if val_depth in observations_all.keys():
            depths = observations_all[val_depth]
            depths[depths == 2 ** 16 - 1] = 0
            depths = depths / 1000

            for i in range(depths.shape[0]):
                frame = frame_name(i)
                dense['depth'][frame] = depths[i]

        ######## EXTRINSICS
        if extrinsics_all is not None and val in extrinsics_all.keys():
            extrinsics = extrinsics_all[val]
            for i in range(extrinsics.shape[0]):
                filename_lowdim = add_key_to_dict(lowdim, filename % (dst, 'lowdim', cam, i, 'npz'))
                lowdim[filename_lowdim]['extrinsics'] = extrinsics[i]

        ######## INTRINSICS
        if intrinsics_all is not None and val in intrinsics_all.keys():
            intrinsics = intrinsics_all[val]
            for i in range(rgbs.shape[0]):                 
                filename_lowdim = add_key_to_dict(lowdim, filename % (dst, 'lowdim', cam, i, 'npz'))
                lowdim[filename_lowdim]['intrinsics'] = intrinsics

        ######## ACTION
        for i in range(action.shape[0]):             
            filename_lowdim = add_key_to_dict(lowdim, filename % (dst, 'lowdim', cam, i, 'npz'))
            lowdim[filename_lowdim]['action'] = {'action': action[i]}
            for obs_key in observations_keys:
                if 'actual' in obs_key or 'desired' in obs_key or 'wrench' in obs_key:
                    lowdim[filename_lowdim]['action'][obs_key] = observations_all[obs_key][i]

        ######## 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='LBM',
            tags=[sim_real,'robotics','tabletop'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=10,
        rgb=dict(extension='jpg'),
        intrinsics=dict(model='pinhole'),
        extrinsics=dict(transform='cam2world',metric=True),
        depth=dict(extension='npz',metric=True,sparse=depth_sparse),
        semantic=None,
        action=dict(format='xyzrot6g'),
        language=dict(task=task,prompt=prompt),
        specific=dict(
            mode=rollout_teleop,
            station=station,
            year=int(year),
            month=int(month),
            day=int(day),
            metadata=metadata,
        ),
    )
    write_json(filename, seq_metadata)

    return dst

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

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

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