# 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_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_binary

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

def get_intrinsics_extrinsics(transforms, i, downsample):

    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

    intrinsics = np.array([
        transforms['fl_x'], transforms['fl_y'],
        transforms['cx'], transforms['cy'],
    ]) / downsample

    return intrinsics, extrinsics

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

def get_sequences(args):
    seqs = crawl(args.src, 'transforms.json', avoid='colmap_480')
    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 = [], {}
    dense_labels = ['rgb','depth']

    ### Get COLMAP data
    try:
        seq_colmap = seq.replace("images_480","colmap_480")
        colmap_path1 = f'{seq_colmap}/colmap'
        colmap_path2 = f'{os.path.dirname(seq_colmap)}/colmap'
        if os.path.exists(f'{colmap_path1}/sparse/0/cameras.bin'):
            cameras_binary, points_binary, images_binary = read_colmap_binary(colmap_path1)
        elif os.path.exists(f'{colmap_path2}/sparse/0/cameras.bin'):
            cameras_binary, points_binary, images_binary = read_colmap_binary(colmap_path2)
        else:
            raise ValueError(f'COLMAP not found {seq}')
        images_binary = {val.name: val for val in images_binary.values()}
        with_colmap = True
    except:
        with_colmap = False

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

    # Camera transformations
    transforms = read_json(f'{seq}/transforms.json')
    file_paths = [os.path.basename(tf['file_path']) for tf in transforms['frames']]
    downsample = {'images_480': 8, 'images_960': 4}[seq.split('/')[4]]

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

        ### Extract COLMAP information
        if with_colmap:
            intrinsics, extrinsics, depth, hw = geometry_from_colmap(
                images_binary[file_path], cameras_binary, points_binary, 
                downsample=downsample, geometry='pinhole')
            ######## DEPTH
            dense['depth'][frame] = depth
        else:
            intrinsics, extrinsics = get_intrinsics_extrinsics(
                transforms, i, downsample)

        ######## RGB
        rgb = np.array(read_image(f'{seq}/images_{downsample}/{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='DL3DV10K',
            tags=['real','static','wild'],
            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=False),
        depth=dict(extension='npz',metric=False,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)

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