# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import numpy as np

from glob import glob

from anydata.utils.read import read_image, read_depth, 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.converters.utils import geometry_from_colmap, read_colmap_text

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

def prepare_transforms(seq, path, file_list):
    transforms = f'{seq}/nerfstudio/{path}.json'
    transforms = read_json(transforms)
    ordered_frames = []
    for i in range(len(file_list)):
        for j in range(len(transforms['frames'])):
            if transforms['frames'][j]['file_path'] == file_list[i]:
                ordered_frames.append(transforms['frames'][j])
                break
    transforms['frames'] = ordered_frames
    return transforms

def get_intrinsics_extrinsics(transforms, i):

    T1 = np.array([[0,1,0,0],[1,0,0,0],[0,0,-1,0],[0,0,0,1]])
    T2 = np.array([[0,-1,0,0],[1,0,0,0],[0,0,1,0],[0,0,0,1]])

    extrinsics = np.array(transforms['frames'][i]['transform_matrix']).reshape(4, 4)
    extrinsics = T1 @ extrinsics @ T1 @ T2

    if transforms['k1'] == 0:  # Undistorted
        intrinsics = np.array([
            transforms['fl_x'], transforms['fl_y'],
            transforms['cx'], transforms['cy'],
        ])
    else: # Distorted
        intrinsics = np.array([
            transforms['fl_x'], transforms['fl_y'],
            transforms['cx'], transforms['cy'],
            transforms['k1'], transforms['k2'],
            transforms['k3'], transforms['k4'],
        ])

    return intrinsics, extrinsics

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

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


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args)

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

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

    dsts = []
    for mode in ['distorted','undistorted']:

        dst = f'{dst_all}/{mode}'
        dsts.append(dst)

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

        ### Get COLMAP data
        if mode == 'distorted':
            try:
                colmap_path = f'{seq}/colmap'
                cameras_binary, points_binary, images_binary = read_colmap_text(colmap_path)
                with_colmap = True
            except:
                with_colmap = False
        else:
            with_colmap = False

        cam = cameras[0]
        dense = {label: dict() for label in dense_labels}

        file_list = read_json(f'{seq}/train_test_lists.json')['train']

        if mode == 'distorted':
            transforms = prepare_transforms(
                seq, 'transforms', file_list)
            intrinsics_model = 'fisheye'
        else:
            transforms_undistorted = prepare_transforms(
                seq, 'transforms_undistorted', file_list)
            intrinsics_model = 'pinhole'

        ############ LOOP OVER CAMERAS
        for i, file_path in enumerate(file_list):
            frame = frame_name(i)

            ### Extract COLMAP information
            if with_colmap:
                val = [val for val in images_binary.values() if val.name == file_path][0]
                intrinsics, extrinsics, depth, hw = geometry_from_colmap(
                    val, cameras_binary, points_binary)
                ######## DEPTH
                dense['depth'][frame] = depth
            else:
                intrinsics, extrinsics = get_intrinsics_extrinsics(transforms, i)

            ######## RGB
            if mode == 'distorted':
                rgb = np.array(read_image(f'{seq}/resized_images/{file_path}'))
            else:
                rgb = np.array(read_image(f'{seq}/resized_undistorted_images/{file_path}'))
            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]['extrinsics'] = extrinsics
            lowdim[filename_lowdim]['intrinsics'] = intrinsics

        ######## 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='ScanNetPP',
                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=intrinsics_model),
            extrinsics=dict(transform='cam2world',metric=False),
            depth=dict(extension='npz',metric=False,sparse=True),
            semantic=None,
            action=None,
            language=None,
            specific=None,
        )
        write_json(filename, seq_metadata)

    return dsts

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

if __name__ == '__main__':
    converter = os.path.basename(__file__)
    run(converter, get_sequences, parse_sequence, process_sequence)

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