# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import torch
import numpy as np

from glob import glob

from anydata.utils.geometry import pose_to_matrix
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

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

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=str, nargs='+')
    parser.add_argument("--num_procs", type=int, default=16)
    parser.add_argument("--local_folder", type=str, default='/data/cv_unified')
    args = parser.parse_args()

    args.src, args.dst = args.path
    args.dst = f'{args.local_folder}/{args.dst}'

    return args

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

def get_sequences(args):
    seqs = crawl(args.src, 'image_left')
    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 = glob(f'{seq}/image*')
    cameras = [c.replace('image_', '').split('/')[-1] for  c in cameras if not c.endswith('.zip')]
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}
    dense_labels = ['rgb','depth','semantic']

    ### Hardcoded intrinsics
    intrinsics = np.array([
        [320.,   0., 320.],
        [  0., 320., 240.],
        [  0.,   0.,   1.],
    ], dtype=np.float32)

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

        ### Get filenames
        filename_rgbs = sorted(glob(f'{seq}/image_{cam}/*.png'))
        filename_depths = sorted(glob(f'{seq}/depth_{cam}/*.npy'))
        filename_semantics = sorted(glob(f'{seq}/seg_{cam}/*.npy'))

        ### Get flow filenames
        filename_flows_masks = sorted(glob(f'{seq}/flow/*.npy'))
        filename_flows = [f for f in filename_flows_masks if 'mask' not in f]
        filename_mask_flows = [f for f in filename_flows_masks if 'mask' in f]
        
        ### Load extrinsics
        extrinsics = np.loadtxt(f'{seq}/pose_{cam}.txt')
        extrinsics = torch.tensor(extrinsics[:, [1, 2, 0, 4, 5, 3, 6]])

        ######## 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    
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            prepare_lowdim(lowdim, dst, cam, frame)

            ######## INTRINSICS + EXTRINSICS
            lowdim[filename_lowdim]['extrinsics'] = pose_to_matrix(extrinsics[i])
            lowdim[filename_lowdim]['intrinsics'] = intrinsics

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

            ######## DEPTH
            depth = read_numpy(filename_depth)
            depth[depth == np.nan] = 0.0
            depth[depth == np.inf] = 0.0
            dense['depth'][frame] = depth

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

            ######## SEGMENTATION
            semantic = read_numpy(filename_semantic)
            dense['semantic'][frame] = semantic

        ######## 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='TartanAir',
            tags=['sim','dynamic','flying','stereo'],
            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=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)

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