# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import numpy as np

from glob import glob
from copy import deepcopy
from collections import namedtuple

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, rotx, roty, rotz, transform_from_rot_trans


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

OxtsPacket = namedtuple(
    'OxtsPacket',
    'lat, lon, alt, ' +
    'roll, pitch, yaw, ' +
    'vn, ve, vf, vl, vu, ' +
    'ax, ay, az, af, al, au, ' +
    'wx, wy, wz, wf, wl, wu, ' +
    'pos_accuracy, vel_accuracy, ' +
    'navstat, numsats, ' +
    'posmode, velmode, orimode'
)

CALIB_FILE = {
    'cam2cam': 'calib_cam_to_cam.txt',
    'velo2cam': 'calib_velo_to_cam.txt',
    'imu2velo': 'calib_imu_to_velo.txt',
}


def get_rgb(filename):
    return np.array(read_image(filename))


def read_npz_depth(file, depth_type):
    """Reads a .npz depth map given a certain depth_type."""
    depth = np.load(file)[depth_type + '_depth'].astype(np.float32)
    return np.expand_dims(depth, axis=2)


def read_png_depth(file):
    """Reads a .png depth map."""
    depth_png = np.array(Image.open(file), dtype=int)
    assert (np.max(depth_png) > 255), 'Wrong .png depth file'
    depth = depth_png.astype(float) / 256.
    depth[depth_png == 0] = -1.
    return np.expand_dims(depth, axis=2)


def get_depth(depth_file):
    """Get the depth map from a file."""
    if depth_file.endswith('.npz'):
        return read_npz_depth(depth_file, 'velodyne')
    elif depth_file.endswith('.png'):
        return read_png_depth(depth_file)
    else:
        raise NotImplementedError(
            'Depth type {} not implemented'.format(self.depth_type))


def get_extrinsics(filename, camera):
    """Gets the pose information from an image file."""

    origin_filename = os.path.dirname(filename) + '/%010d.txt' % 0
    if not os.path.exists(origin_filename):
        return None

    origin_oxts_data = np.loadtxt(origin_filename, delimiter=' ', skiprows=0)
    lat = origin_oxts_data[0]
    scale = np.cos(lat * np.pi / 180.)
    origin_R, origin_t = pose_from_oxts_packet(origin_oxts_data, scale)
    origin_pose = transform_from_rot_trans(origin_R, origin_t)

    oxts_data = np.loadtxt(filename, delimiter=' ', skiprows=0)
    R, t = pose_from_oxts_packet(oxts_data, scale)
    pose = transform_from_rot_trans(R, t)
    imu2cam = get_imu2cam_transform(filename)
    odo_pose = (imu2cam @ np.linalg.inv(origin_pose) @
                pose @ np.linalg.inv(imu2cam)).astype(np.float32)

    if camera == 'right':
        odo_pose[0, -1] += 0.5407

    return odo_pose


def pose_from_oxts_packet(raw_data, scale):
    """Helper method to compute a SE(3) pose matrix from an OXTS packet"""

    packet = OxtsPacket(*raw_data)
    er = 6378137.  # earth radius (approx.) in meters

    # Use a Mercator projection to get the translation vector
    tx = scale * packet.lon * np.pi * er / 180.
    ty = scale * er * \
        np.log(np.tan((90. + packet.lat) * np.pi / 360.))
    tz = packet.alt
    t = np.array([tx, ty, tz])

    # Use the Euler angles to get the rotation matrix
    Rx = rotx(packet.roll)
    Ry = roty(packet.pitch)
    Rz = rotz(packet.yaw)
    R = Rz.dot(Ry.dot(Rx))

    # Combine the translation and rotation into a homogeneous transform
    return R, t


def read_calib_file(filepath):
    """Read in a calibration file and parse into a dictionary"""
    data = {}
    with open(filepath, 'r') as f:
        for line in f.readlines():
            key, value = line.split(':', 1)
            try:
                data[key] = np.array([float(x) for x in value.split()])
            except ValueError:
                pass
    return data


def get_imu2cam_transform(filename):
    """Gets the transformation between IMU an camera from an image file"""
    parent_folder = '/'.join(filename.split('/')[:-4])

    cam2cam = read_calib_file(os.path.join(parent_folder, CALIB_FILE['cam2cam']))
    imu2velo = read_calib_file(os.path.join(parent_folder, CALIB_FILE['imu2velo']))
    velo2cam = read_calib_file(os.path.join(parent_folder, CALIB_FILE['velo2cam']))

    velo2cam_mat = transform_from_rot_trans(velo2cam['R'], velo2cam['T'])
    imu2velo_mat = transform_from_rot_trans(imu2velo['R'], imu2velo['T'])
    cam_2rect_mat = transform_from_rot_trans(cam2cam['R_rect_00'], np.zeros(3))

    imu2cam = cam_2rect_mat @ velo2cam_mat @ imu2velo_mat
    return imu2cam


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

def get_sequences(args):
    seqs = crawl(args.src, 'oxts')
    seqs = [os.path.dirname(seq) for seq in seqs]
    return seqs


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args, remove=[3])

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

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

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

    cam_map = {cameras[i]: camera_names[i] for i in range(len(cameras))}

    single_intrinsics = np.array([
        [0.58, 0.00, 0.5],
        [0.00, 1.92, 0.5],
        [0.00, 0.00, 1.0]], dtype=np.float32)

    filename_oxts = sorted(glob(f'{seq}/oxts/data/*.txt'))

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

        ### Get filenames
        filename_rgbs = sorted(glob(f'{seq}/{cam}/data/0*.png'))
        filename_depths = sorted(glob(f'{seq}/proj_depth/velodyne/{cam}/0*.npz'))

        ### Change camera name
        cam = cam_map[cam]

        ######## 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
        for i, filename_oxt in enumerate(filename_oxts):
            frame = frame_name(i)

            extrinsics = get_extrinsics(filename_oxt, cam)

            intrinsics = deepcopy(single_intrinsics)
            intrinsics[0, :] *= rgb.shape[1]
            intrinsics[1, :] *= rgb.shape[0]

            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['intrinsics'] = intrinsics
            if extrinsics is not None:
                lowdim[filename_lowdim]['extrinsics'] = extrinsics

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

            ### DEPTH
            depth = get_depth(filename_depth)[..., 0]
            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='KITTI',
            tags=['real','dynamic','driving'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=camera_names,
        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 dst

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

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

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