# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import cv2
import numpy as np

from glob import glob

from anydata.utils.geometry import invert_extrinsics
from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth, read_json
from anydata.utils.write import write_image, write_npz, write_json
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, crawl

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

def convert_extrinsics(extrinsics):
    flip_yz = np.eye(4)
    flip_yz[1, 1] = -1
    flip_yz[2, 2] = -1
    extrinsics = np.matmul(extrinsics, flip_yz)
    return extrinsics

def get_extrinsics(filename, transforms):
    filename = filename.replace('images', 'poses').replace('.png', '.json')
    extrinsics = np.array(read_json(filename)['pose'])
    extrinsics = convert_extrinsics(extrinsics)
    return extrinsics

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

def get_sequences(args):
    seqs = crawl(args.src, 'transforms.json')
    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 = ['0']
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}

    ### Get filenames
    filename_rgbs = sorted(glob(f'{seq}/images/*.png'))
    filename_poses = sorted(glob(f'{seq}/poses/*.json'))

    ### Read transforms for the sequence
    transforms = read_json(f'{seq}/transforms.json')

    ### Get intrinsics
    focal, cx, cy = transforms["fl_x"], transforms["cx"], transforms["cy"]
    intrinsics = np.array([[focal, 0, cx], [0, focal, cy], [0, 0, 1]])

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

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

            ### Get extrinsics
            extrinsics = get_extrinsics(filename_rgb, transforms)

            ######## INTRINSICS + EXTRINSICS
            if 'intrinsics' not in labels: labels.append('intrinsics')
            if 'extrinsics' not in labels: labels.append('extrinsics')
            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_rgbs) > 0:
            resolution[cam]['rgb'] = rgb.shape[:2]
            num_frames[cam]['rgb'] = len(filename_rgbs)

    ######## 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='HM3D',
            tags=['sim','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=None,
        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)

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