# 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_yaml, read_image, read_depth
from anydata.utils.write import write_image, write_npz, write_json, write_jpg_from_mp4
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, crawl

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

def get_depth(filename):
    depth = read_depth(filename, div=1000)
    return depth
    
def get_mask(filename):
    return read_image(filename, '1')
    
#######################################################

def get_sequences(args):
    seqs = crawl(args.src, 'align_rgb')
    return seqs


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

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

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

    ### Initialize lists and dicts
    cameras = ['head']
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}

    ############ LOOP OVER CAMERAS
    for cam in cameras:

        vid = f'{seq}/image.mp4'
        folder_rgb_out = f'{dst}/rgb/{cam}'
        resolution_rgb, num_rgb_frames = write_jpg_from_mp4(folder_rgb_out, vid)

        split = seq.split('/')
        intrinsics = np.load(f'{args.src}/camera_params/{split[-8]}/intrin.npy')
        print(intrinsics)
        # intrinsics = f'{seq}/'


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

        #     ######## RGB
        #     if 'rgb' not in labels: labels.append('rgb')
        #     rgb = np.array(read_image(filename_rgb).resize((640, 480)))
        #     filename_rgb_out = f'{dst}/rgb/{cam}/{frame}.jpg'
        #     write_image(filename_rgb_out, rgb)

        #     ######## LOWDIM RGB             
        #     filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
        #     lowdim[filename_lowdim]['camera'] = cam
        #     lowdim[filename_lowdim]['timestep'] = int(frame)

        # if len(filename_rgbs) > 0:
        #     resolution[cam]['rgb'] = rgb.shape[:2]
        #     num_frames[cam]['rgb'] = len(filename_rgbs)

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

        #     ######## INTRINSICS + EXTRINSICS
        #     if 'intrinsics' not in labels: labels.append('intrinsics')
        #     if 'extrinsics' not in labels: labels.append('extrinsics')
        #     extrinsics = np.genfromtxt(filename_extrinsic).astype(np.float32).reshape((4, 4))
        #     filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
        #     lowdim[filename_lowdim]['extrinsics'] = extrinsics
        #     lowdim[filename_lowdim]['intrinsics'] = intrinsics

        # if len(filename_extrinsics) > 0:
        #     num_frames[cam]['extrinsics'] = len(filename_extrinsics)

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

        #     ######## DEPTH
        #     if 'depth' not in labels: labels.append('depth')
        #     depth = cv2.imread(filename_depth, cv2.IMREAD_ANYDEPTH).astype(np.float32) / 1000.0
        #     filename_depth_out = f'{dst}/depth/{cam}/{frame}.npz'
        #     write_npz(filename_depth_out, dict(depth=depth))

        # if len(filename_depths) > 0:
        #     resolution[cam]['depth'] = depth.shape[:2]
        #     num_frames[cam]['depth'] = len(filename_depths)

    ######## WRITE LOWDIM
    for key, val in lowdim.items():
        write_npz(key, val)

############ METADATA 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='ScanNet',
            tags=['real','static','indoor'],
            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=True),
        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)

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