# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import torch
import numpy as np
import pandas as pd

from glob import glob

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.converters.utils import run, add_key_to_dict, fill_metadata, read_colmap_binary, geometry_from_colmap, parse_dst_seq, frame_name, crawl
from anydata.geometry.camera_utils import multiply_extrinsics, invert_extrinsics

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

CAMERAS = [
    'front_pvm', 
    'front_tele', 
    'front_wide', 
    'left_side_forward', 
    'left_side_pvm', 
    'left_side_rearward', 
    'rear_pvm', 
    'rear_tele', 
    'right_side_forward', 
    'right_side_pvm', 
    'right_side_rearward'
]

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

def get_sequences(args):
    seqs = crawl(args.src, '*.parquet')
    # seqs = [s[:-len('.parquet')] for s in seqs]
    print(seqs)
    return seqs


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

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

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

    df = pd.read_parquet(seq)

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

    dsts, scene = [], None
    for i in range(len(df)):

        scene_id = df.loc[i]['scene_id']
        frame_idx = int(df.loc[i]['frame_idx'])

        if scene != scene_id:
            dst = f'{dst_all}/{scene_id}'
            dsts.append(dst)
            scene = scene_id
            cnt = 0

        smooth_from_base = df.loc[i][f'smooth_from_base'] 
        smooth_from_base = np.stack(smooth_from_base, 0)
        smooth_from_base = torch.tensor(smooth_from_base).unsqueeze(0).float()

        pcl_path = f'{dst}/pointcloud/%010d' % i
        os.makedirs(f'{dst}/pointcloud', exist_ok=True)

        os.system(f'aws s3 cp {df.loc[i]["x"]} {pcl_path}/x.npy --profile tri-e2e-dev')
        os.system(f'aws s3 cp {df.loc[i]["y"]} {pcl_path}/y.npy --profile tri-e2e-dev')
        os.system(f'aws s3 cp {df.loc[i]["z"]} {pcl_path}/z.npy --profile tri-e2e-dev')

        x = torch.tensor(np.load(f'{pcl_path}/x.npy'))
        y = torch.tensor(np.load(f'{pcl_path}/y.npy'))
        z = torch.tensor(np.load(f'{pcl_path}/z.npy'))
        # remove_path(pcl_path)

        pointcloud = torch.stack([x, y, z], 0).unsqueeze(0).float()
        pointcloud = multiply_extrinsics(pointcloud, smooth_from_base)    

        cams = []
        for k, cam in enumerate(cameras):

            print("#################################", i, k, cam)

            rgb = df.loc[i][f'{cam}_image']            
            width = int(df.loc[i][f'{cam}_calibration_camera_width'])
            height = int(df.loc[i][f'{cam}_calibration_camera_height']) 
            focal_length = df.loc[i][f'{cam}_calibration_focal_length'] 
            optical_center = df.loc[i][f'{cam}_calibration_optical_center'] 
            distortion_coefs = df.loc[i][f'{cam}_calibration_distortion_coefs'] 
            camera_from_vehicle = df.loc[i][f'{cam}_calibration_camera_from_vehicle_matrix'] 

            resolution[cam] = {'rgb': (height, width)}
            num_frames[cam] = {'rgb': frame_idx + 1}

            command = f'aws s3 cp {rgb} {dst}/rgb/{cam}/%010d.jpg --profile tri-e2e-dev' % frame_idx
            os.system(command)

            intrinsics = np.array([
                focal_length['x'], focal_length['y'], 
                optical_center['x'], optical_center['y'], 
                distortion_coefs['k1'], distortion_coefs['k2'], 
                distortion_coefs['k3'], distortion_coefs['k4'],
            ])

            extrinsics = np.stack(camera_from_vehicle, 0)
            extrinsics = torch.tensor(extrinsics).unsqueeze(0).float()
            extrinsics = invert_extrinsics(extrinsics)
            extrinsics = smooth_from_base @ extrinsics

            from anydata.geometry.camera import Camera
            camera = Camera(
                K=torch.tensor(intrinsics).unsqueeze(0).float(), 
                Tcw=extrinsics, 
                hw=(height,width), 
                geometry='distorted',
            )
            cams.append(camera)

            depth = camera.project_pointcloud(pointcloud, from_world=True)
            points = camera.reconstruct_depth_map(depth, to_world=True)
            depth = depth[0, 0].numpy()

            lowdim = {
                'extrinsics': extrinsics[0].numpy(),
                'intrinsics': intrinsics,
                'timestep': frame_idx,
                'camera': cam,
            }
            
            os.makedirs(f'{dst}/lowdim/{cam}', exist_ok=True)
            np.savez_compressed(f'{dst}/lowdim/{cam}/%010d.npz' % frame_idx, **lowdim)

            os.makedirs(f'{dst}/depth/{cam}', exist_ok=True)
            np.savez_compressed(f'{dst}/depth/{cam}/%010d.npz' % frame_idx, depth=depth)

    ############ METADATA 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(name='PhysicalAI', tags=['real']),
        labels=["rgb", "intrinsics", "extrinsics", "depth"],
        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 dsts

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

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

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