# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import csv
import cv2
import numpy as np
import torch 

from glob import glob
from PIL import Image

from anydata.utils.read import read_image
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.utils.geometry import invert_extrinsics

from anydata.geometry.camera import Camera
from anydata.geometry.depth import calculate_normals, calc_dot_prod

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

from projectaria_tools.projects import ase
from projectaria_tools.core import calibration
from scipy.spatial.transform import Rotation as R

device = ase.get_ase_rgb_calibration()
pose_device = device.get_transform_device_camera().to_matrix()
pinhole = calibration.get_linear_camera_calibration(
    device.get_image_size()[0],
    device.get_image_size()[1],
    device.get_focal_lengths()[0])    

cx, cy = pinhole.get_principal_point()
fx, fy = pinhole.get_focal_lengths()

calib = np.array([
    [ fy, 0.0,  cy],
    [0.0,  fx,  cx], 
    [0.0, 0.0, 1.0]
])

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

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

def get_rgb(filename):
    rgb = read_image(filename)
    rgb = calibration.distort_by_calibration(rgb, pinhole, device)
    rgb = Image.fromarray(rgb)
    return np.array(rgb.rotate(-90))

def get_intrinsics():
    return calib

def get_extrinsics(trajectory, i):

    line = trajectory[i+1]    
    translation = np.array([float(p) for p in line[3:6]])
    quat_xyzw = np.array([float(o) for o in line[6:10]])
    rotation = R.from_quat(quat_xyzw).as_matrix()

    pose = np.eye(4)
    pose[:3, :3] = rotation
    pose[:3, -1] = translation
    pose = pose @ pose_device

    pose = invert_extrinsics(pose)
    pose = pose_fix @ pose @ pose_fix
    pose = invert_extrinsics(pose)
    return pose



def get_depth(filename):

    depth = Image.open(filename)
    depth = np.array(depth).astype(np.float32)
    depth = calibration.distort_by_calibration(depth, pinhole, device) / 1000

    hw = device.get_image_size()
    K = torch.tensor(calib).unsqueeze(0).float()
    cam = Camera(hw=hw, K=K.cuda())

    depth = torch.tensor(depth).unsqueeze(0).unsqueeze(0).cuda()
    depth = cam.e2z(depth).cpu()

    depth = depth.squeeze().numpy()
    return np.rot90(depth, k=3).copy()

def get_mask_depth(depth):

    hw = device.get_image_size()
    K = torch.tensor(calib).unsqueeze(0).float()
    cam = Camera(hw=hw, K=K)

    depth2 = torch.tensor(depth).unsqueeze(0).unsqueeze(0)

    normals = calculate_normals(depth2, camera=cam)
    points = cam.reconstruct_depth_map(depth2)
    dotprod = calc_dot_prod(points, normals)
    dotprod[torch.isnan(dotprod)] = 0
    invalid = dotprod[:, 0].abs() < 0.1

    mask = torch.ones_like(depth2)
    mask[invalid.unsqueeze(1)] = 0
    mask = mask.squeeze().numpy()

    return (mask * 255).astype(np.uint8)
        
            
#######################################################

def get_sequences(args):
    seqs = crawl(args.src, 'trajectory.csv')
    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 filenames
    filename_rgbs = sorted(glob(f'{seq}/rgb/*.jpg'))
    filename_depths = sorted(glob(f'{seq}/depth/*.png'))

    ### Get trajectories
    trajectory = f'{seq}/trajectory.csv'
    with open(trajectory) as file:
        trajectory = [f for f in csv.reader(file)]

    ############ LOOP OVER CAMERAS
    for cam in cameras:
        dense = {label: dict() for label in dense_labels}

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

            ######## RGB
            rgb = get_rgb(filename_rgb)
            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'] = get_extrinsics(trajectory, i)
            lowdim[filename_lowdim]['intrinsics'] = get_intrinsics()

        ######## DEPTH FILENAMES
        for i, filename_depth in enumerate(filename_depths):
            frame = frame_name(i)

            ######## DEPTH
            depth = get_depth(filename_depth)
            dense['depth'][frame] = depth

        ######## 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='AriaSynthetic',
            tags=['sim','static','indoors'],
            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=dict(extension='npz',metric=True,sparse=False),
        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)

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