# Copyright 2026 Toyota Research Institute.  All rights reserved.

from tqdm import tqdm
from glob import glob

import os
import torch
import argparse
import numpy as np
from PIL import Image

from anydata.converters.utils import add_key_to_dict, fill_metadata, get_splits, geometry_from_colmap

from anydata.utils.write import write_image, write_npz, write_json
from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth
from anydata.utils.colmap import read_cameras_binary, read_images_binary, read_points3d_binary, qvec2rotmat

from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth
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, invert_extrinsics

from anydata.geometry.camera_utils import multiply_extrinsics

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

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_depth(filename):
    depth = read_depth(filename, div=1000)
    return depth
    
def get_mask(filename):
    return read_image(filename, '1')

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

def get_sequences(args):
    seqs = glob(f'{args.src}/**/transforms.json', recursive=True)
    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):

    labels = []
    cameras = ['left','right']
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    lowdim = dict()

    ### Get COLMAP data
    from anydata.converters.utils import read_colmap_binary
    cameras_binary, points_binary, images_binary = read_colmap_binary(seq)

    from anydata.utils.read import read_json
    transforms = read_json(f'{seq}/transforms.json')
    undistort = transforms['undistort_camera_model']
    frames = transforms['frames']

    intrinsics = np.array(undistort['intrinsic'])
    intrinsics[:-1] /= 2

    import open3d as o3d

    # Read .ply file
    input_file = f"{seq}/raw_pcd.ply"
    pcd = o3d.io.read_point_cloud(input_file) # Read the point cloud
    points = np.asarray(pcd.points) 
    points = torch.tensor(points).permute(1, 0).unsqueeze(0).float()

    # transform = np.array([
    #     [ 1.0,  0.0,  0.0,  0.0],
    #     [ 0.0,  1.0,  0.0,  0.0],
    #     [ 0.0,  0.0, -1.0,  0.0],
    #     [ 0.0,  0.0,  0.0,  1.0],
    # ])

    transform = np.array([
        [ 0.0,  1.0,  0.0,  0.0],
        [ 1.0,  0.0,  0.0,  0.0],
        [ 0.0,  0.0,  -1.0,  0.0],
        [ 0.0,  0.0,  0.0,  1.0],
    ])

    transform = torch.tensor(transform).unsqueeze(0).float()

    points = multiply_extrinsics(points,transform)

    # # Visualize the point cloud within open3d
    # o3d.visualization.draw_geometries([pcd]) 

    ############ LOOP OVER CAMERAS
    for i, (key, val) in enumerate(images_binary.items()):
        frame = frame_name(val.name, -1)

        ### Extract COLMAP information
        intrinsics, extrinsics, depth, hw = geometry_from_colmap(
            val, cameras_binary, points_binary)

        ######## RGB
        if 'rgb' not in labels: labels.append('rgb')
        rgb = np.array(read_image(f'{seq}/images/{val.name}'))
        filename_rgb_out = f'{dst}/rgb/{cam}/{frame}.jpg'
        write_image(filename_rgb_out, rgb)

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

        ######## DEPTH
        if 'depth' not in labels: labels.append('depth')
        filename_depth_out = f'{dst}/depth/{cam}/{frame}.npz'
        write_npz(filename_depth_out, dict(depth=depth))

        ######## 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(images_binary) > 0:
        resolution[cam]['rgb'] = rgb.shape[:2]
        num_frames[cam]['rgb'] = len(images_binary)
        resolution[cam]['depth'] = depth.shape[:2]
        num_frames[cam]['depth'] = len(images_binary)


    # for cam in cameras:

    #     filename_rgbs = sorted(glob(f'{seq}/images/{cam}*'))

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

    #         name = filename_rgb.split('/')[-1].replace('_','\\').replace('.png', '.jpg')
    #         transform_matrix = [s for s in frames if s['file_path'] == name][0]['transform_matrix']
    #         extrinsics = np.array(transform_matrix) 
    #         extrinsics = invert_extrinsics(extrinsics)

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

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

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

    #         from anydata.geometry.camera import Camera
    #         camera = Camera(
    #             K=torch.tensor(intrinsics).unsqueeze(0).float(),
    #             Tcw=torch.tensor(extrinsics).unsqueeze(0).float(),
    #             hw=rgb.shape[:2],
    #         )
            
    #         depth = camera.project_pointcloud(points)[0,0].numpy()

    #         ### DEPTH
    #         if 'depth' not in labels: labels.append('depth')
    #         filename_depth_out = f'{dst}/depth/{cam}/{frame}.npz'
    #         write_npz(filename_depth_out, dict(depth=depth))

    #         if i > 1:
    #             break

    #     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='Wanderland',
            tags=['real'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=fps,
        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)

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