# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import json
import numpy as np

from glob import glob

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

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

def get_rgb(filename):
    return np.array(read_image(filename))


def get_depth(filename, key):
    return read_numpy(filename, key=key)


def get_intrinsic(filename):
    return np.load(filename)


def get_extrinsic(filename):
    return invert_extrinsics(np.load(filename))


def get_semantic(filename):
    return np.array(read_image(filename, mode=''))[..., 0]


def get_optflow(filename):
    optflow = np.array(read_image(filename, mode='RGBA'), dtype=np.uint16)
    # Convert to uv motion
    dx_i = optflow[..., 0] + optflow[..., 1] * 256
    dy_i = optflow[..., 2] + optflow[..., 3] * 256
    dx = ((dx_i / 65535.0) * 2.0 - 1.0)
    dy = ((dy_i / 65535.0) * 2.0 - 1.0)
    # Return stacked array
    return np.stack((dx, dy), 2)


def get_bbox2d(filename):
    try:
        data = read_json(filename)['annotations'] 
        if len(data) == 0:
            return np.zeros((0, 7))
    except:
        return np.zeros((0, 7))

    bboxes2d = []
    for d in data:
        class_id = d['class_id']
        instance_id = d['instance_id']
        visible = json.loads(d['attributes']['user_data'])['visibility']
        x, y, w, h = d['box']['x'], d['box']['y'], d['box']['w'], d['box']['h']
        bboxes2d.append([x, y, x + w, y + h, visible, class_id, instance_id])
    return np.array(bboxes2d)


def get_bbox3d(filename):
    try:
        data = read_json(filename)['annotations']
        if len(data) == 0:
            return np.zeros((0, 13))
    except:
        return np.zeros((0, 13))

    bboxes3d = []
    for d in data:

        h = d['box']['height']
        w = d['box']['width']
        l = d['box']['length']

        class_id = d['class_id']
        instance_id = d['instance_id']
        visible = d['num_points']

        tvec = d['box']['pose']['translation']
        qvec = d['box']['pose']['rotation']

        dims = [h, w, l]
        tvec = [tvec['x'], tvec['y'], tvec['z']]
        qvec = [qvec['qw'], qvec['qx'], qvec['qy'], qvec['qz']]
        prop = [visible, class_id, instance_id]

        bboxes3d.append(dims + tvec + qvec + prop)

    return np.array(bboxes3d)


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

def get_sequences(args):
    seqs = crawl(args.src, 'scene*.json')
    seqs = [os.path.dirname(seq) for seq in seqs]
    seqs = list(set(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 = [os.path.basename(cam) for cam in sorted(glob(f'{seq}/rgb/*'))]
    cameras = [c for c in cameras if c not in ['camera_04'] and 'virtual' not in c] # Invalid older PD cameras
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}
    dense_labels = ['rgb','depth','optflow_fwd','optflow_bwd','semantic']

    # Sim/real differences
    real = not seq.split('/')[-2].startswith('PD')
    depth_prefix = 'projected/depth/lidar' if real else 'depth'
    depth_key = 'depth' if real else 'data'
    sim_real = 'real' if real else 'sim'
    depth_sparse = True if real else False

    ############ 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}/*.png'))
        if len(filename_rgbs) == 0:
            filename_rgbs = sorted(glob(f'{seq}/rgb/{cam}/*.jpg'))
        filename_depths = sorted(glob(f'{seq}/{depth_prefix}/{cam}/*.npz'))
        filename_intrinsics = sorted(glob(f'{seq}/intrinsics/{cam}/*.npy'))
        filename_extrinsics = sorted(glob(f'{seq}/pose/{cam}/*.npy'))
        
        filename_optflow_fwds = sorted(glob(f'{seq}/motion_vectors_2d/{cam}/*.png'))
        filename_optflow_bwds = sorted(glob(f'{seq}/back_motion_vectors_2d/{cam}/*.png'))

        filename_semantics = [] if real else sorted(glob(f'{seq}/semantic_segmentation_2d/{cam}/*.png'))

        filename_bbox2ds = sorted(glob(f'{seq}/bounding_box_2d/{cam}/*.json'))
        filename_bbox3ds = sorted(glob(f'{seq}/bounding_box_3d/{cam}/*.json'))

        ######## OPTFLOW_FWD FILENAMES
        for i, filename_optflow_fwd in enumerate(filename_optflow_fwds):
            frame = frame_name(i)

            ######## OPTFLOW_FWD
            optflow_fwd = get_optflow(filename_optflow_fwd)
            dense['optflow_fwd'][frame] = optflow_fwd

        if len(filename_optflow_fwds) > 0:  # add last frame as invalid
            frame = frame_name(len(filename_optflow_fwds))
            dense['optflow_fwd'][frame] = -1 * np.ones_like(optflow_fwd)

        ######## OPTFLOW_BWD FILENAMES
        for i, filename_optflow_bwd in enumerate(filename_optflow_bwds):
            frame = frame_name(i+1)

            ######## OPTFLOW_BWD
            optflow_bwd = get_optflow(filename_optflow_bwd)
            dense['optflow_bwd'][frame] = optflow_bwd

        if len(filename_optflow_bwds) > 0:  # Add first frame as invalid
            frame = frame_name(0)
            dense['optflow_bwd'][frame] = -1 * np.ones_like(optflow_bwd)

        ######## BBOX2D FILENAMES
        for i, filename_bbox2d in enumerate(list(filename_bbox2ds)):
            frame = frame_name(i)

            ######## BBOX2D
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['bbox2d'] = get_bbox2d(filename_bbox2d)

        ######## BBOX3D FILENAMES
        for i, filename_bbox3d in enumerate(list(filename_bbox3ds)):
            frame = frame_name(i)

            ######## BBOX3D
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['bbox3d'] = get_bbox3d(filename_bbox3d)

        ######## SEMANTIC FILENAMES
        for i, filename_semantic in enumerate(filename_semantics):
            frame = frame_name(i)

            ######## SEMANTIC
            semantic = get_semantic(filename_semantic)
            dense['semantic'][frame] = semantic

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

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

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

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

            ### DEPTH
            depth = get_depth(filename_depth, key=depth_key)
            dense['depth'][frame] = depth

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

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

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

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

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

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

    ######## ONTOLOGY
    if len(filename_semantics) > 0:
        ontology = {}
        original_ontology = read_json(glob(f'{seq}/ontology/*.json')[0])['items']
        for o in original_ontology:
            ontology[o['id']] = dict(
                name=o['name'],
                color=[o['color']['r'], o['color']['g'], o['color']['b']],
                thing_or_stuff='thing' if o['isthing'] else 'stuff',
            )
        semantic = dict(ontology=ontology)
    else:
        semantic = None

############ METADATA 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name=os.path.basename(args.src),
            tags=[sim_real,'dynamic','driving'],
            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=semantic, 
        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)

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