# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import cv2
import numpy as np

from glob import glob

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

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

def get_sequences(args):
    seqs = glob(f'{args.src}/*')
    seqs = [s for s in seqs if not s.endswith('.json')]
    return seqs


def parse_sequence(seq, args):
    name = os.path.basename(args.src)
    if   name == 'NuScenes': remove = [0]
    elif name == 'PD4D':     remove = [0]
    else:                    remove = []
    return parse_dst_seq(seq, args, separate='__', remove=remove)

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

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

    ### Initialize lists and dicts
    cameras = sorted(glob(f'{seq}/rgb/*'))
    cameras = [os.path.basename(c) for c in cameras]
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}
    dense_labels = ['rgb','depth']

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

        ### Get filenames
        filename_rgbs = sorted(glob(f'{seq}/rgb/{cam}/*.jpg'))
        filename_depths = sorted(glob(f'{seq}/depth/{cam}/*.npz'))
        filename_intrinsics = sorted(glob(f'{seq}/intrinsics/{cam}/*.npy'))
        filename_extrinsics = sorted(glob(f'{seq}/pose/{cam}/*.npy'))
        filename_lowdims = sorted(glob(f'{seq}/lowdim/{cam}/*.npz'))

        if len(filename_depths) == 0:
            filename_depths = sorted(glob(f'{seq}/depth/zbuffer/{cam}/*.npz'))

        ######## RGB FILENAMES
        for i, filename_rgb in enumerate(filename_rgbs):
            frame = frame_name(i)

            ######## RGB
            rgb = np.array(read_image(filename_rgb))
            dense['rgb'][frame] = rgb

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

        ######## EXTRINSICS
        for i, filename_extrinsic in enumerate(filename_extrinsics):
            frame = frame_name(i)

            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['extrinsics'] = invert_extrinsics(read_numpy(filename_extrinsic))

        ######## INTRINSICS
        for i, filename_intrinsic in enumerate(filename_intrinsics):
            frame = frame_name(i)

            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['intrinsics'] = read_numpy(filename_intrinsic)[:3, :3]

        ######## LOWDIM
        for i, filename_ld in enumerate(filename_lowdims):
            frame = frame_name(i)

            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            ld = read_numpy(filename_ld)
            lowdim[filename_lowdim]['intrinsics'] = ld['intrinsics'][:3, :3]
            lowdim[filename_lowdim]['extrinsics'] = ld['pose']

        ######## DEPTH FILENAMES
        for i, filename_depth in enumerate(filename_depths):
            frame = frame_name(i)

            ######## DEPTH
            depth = np.load(filename_depth)['data']
            dense['depth'][frame] = depth

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

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

############ METADATA 
    name = os.path.basename(args.src)
    if name.startswith('PD'):  # Parallel Domain
        tags = ['sim','dynamic','driving']
        depth_sparse, metric = False, True
    elif name.startswith('Kubric'):  # Kubric datasets
        tags = ['sim','dynamic','indoors','inward']
        depth_sparse, metric = False, False
    else:  # Driving datasets
        tags = ['real','dynamic','driving']
        depth_sparse, metric = True, True
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name=name,
            tags=tags,
            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=metric),
        depth=dict(extension='npz',metric=metric,sparse=depth_sparse),
        semantic=None,
        action=None,
        language=None,
        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)

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