# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import json
import numpy as np
import cv2

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
from anydata.utils.colmap import qvec2rotmat

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

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


def get_depth(filename, **kwargs):
    """Get depth from filename"""
    return cv2.imread(filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) / 100.


def get_intrinsics(intrinsics_all, i, c):
    intrinsics = intrinsics_all[2 * i + c]
    fx, fy, cx, cy = intrinsics[2:]
    return np.array([
        [ fx, 0.0,  cx],
        [0.0,  fy,  cy],
        [0.0, 0.0, 1.0],
    ])


def get_extrinsics(extrinsics_all, i, c):
    extrinsics = extrinsics_all[2 * i + c]
    extrinsics = np.array(extrinsics[2:]).reshape(4, 4)
    extrinsics = invert_extrinsics(extrinsics)
    return extrinsics


def get_optflow(filename):
    """Get optical flow from filename"""
    # Get optical flow
    optical_flow = cv2.imread(filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
    h, w = optical_flow.shape[:2]
    # Get invalid optical flow pixels
    invalid = optical_flow[..., 0] == 0
    # Normalize and scale optical flow values
    optical_flow = 2.0 / (2 ** 16 - 1.0) * optical_flow[..., 2:0:-1].astype('f4') - 1.
    # Remove invalid pixels
    optical_flow[invalid] = 0
    return optical_flow

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

def get_sequences(args):
    seqs = glob(f'{args.src}/**/rgb/Camera_0', recursive=True)
    seqs = [os.path.dirname(seq) for seq in seqs]
    seqs = [os.path.dirname(seq) for seq in seqs]
    return seqs


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args, remove=[0,3])

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

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

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

    cam_map = {cameras[i]: camera_names[i] for i in range(len(cameras))}

    from anydata.utils.read import read_txt
    tmpl = os.path.dirname(seq.replace('/rgb/', '/textgt/'))
    intrinsics_all = read_txt(f'{tmpl}/intrinsic.txt')[1:]
    for i in range(len(intrinsics_all)):
        intrinsics_all[i] = [float(v) for v in intrinsics_all[i].split(' ')]
    extrinsics_all = read_txt(f'{tmpl}/extrinsic.txt')[1:]
    for i in range(len(extrinsics_all)):
        extrinsics_all[i] = [float(v) for v in extrinsics_all[i].split(' ')]

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

        tmpl = seq.replace('/rgb/', '/%s/')
        tmpl = f'{tmpl}/%s/{cam}'

        ### Get filenames
        filename_rgbs = sorted(glob(f'{tmpl}/*.jpg' % ('rgb','rgb')))
        filename_depths = sorted(glob(f'{tmpl}/*.png' % ('depth','depth')))

        ### Change camera name
        cam = cam_map[cam]

        ######## 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)

            ######## INTRINSICS + EXTRINSICS             
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['intrinsics'] = get_intrinsics(intrinsics_all, i, c)
            lowdim[filename_lowdim]['extrinsics'] = get_extrinsics(extrinsics_all, i, c)

            ######## OPTFLOW_FWD
            filename_optflow_fwd = filename_rgb.replace('/rgb/', '/forwardFlow/').replace('.jpg', '.png')
            filename_optflow_fwd = filename_optflow_fwd.replace('/rgb_', '/flow_')
            if os.path.exists(filename_optflow_fwd):
                optflow_fwd = get_optflow(filename_optflow_fwd)
                dense['optflow_fwd'][frame] = optflow_fwd
            else:
                dense['optflow_fwd'][frame] = -1 * np.ones((rgb.shape[0], rgb.shape[1], 2))


            ######## OPTFLOW_BVWD
            filename_optflow_bwd = filename_rgb.replace('/rgb/', '/backwardFlow/').replace('.jpg', '.png')
            filename_optflow_bwd = filename_optflow_bwd.replace('/rgb_', '/backwardFlow_')
            if os.path.exists(filename_optflow_bwd):
                optflow_bwd = get_optflow(filename_optflow_bwd)
                dense['optflow_bwd'][frame] = optflow_bwd
            else:
                dense['optflow_bwd'][frame] = -1 * np.ones((rgb.shape[0], rgb.shape[1], 2))

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

            ### DEPTH
            depth = get_depth(filename_depth, key='depth')
            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 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='VKITTI2',
            tags=['sim','dynamic','driving'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=camera_names,
        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=False),
        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)

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