# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import gzip
import json
import torch
import numpy as np

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

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

def opencv_from_cameras_projection(R, T, focal, p0, image_size):
    R = torch.from_numpy(R)[None, :, :]
    T = torch.from_numpy(T)[None, :]
    focal = torch.from_numpy(focal)[None, :]
    p0 = torch.from_numpy(p0)[None, :]
    image_size = torch.from_numpy(image_size)[None, :]

    R_pytorch3d = R.clone()
    T_pytorch3d = T.clone()
    focal_pytorch3d = focal
    p0_pytorch3d = p0
    T_pytorch3d[:, :2] *= -1
    R_pytorch3d[:, :, :2] *= -1
    tvec = T_pytorch3d
    R = R_pytorch3d.permute(0, 2, 1)

    # Retype the image_size correctly and flip to width, height.
    image_size_wh = image_size.to(R).flip(dims=(1,))

    # NDC to screen conversion.
    scale = image_size_wh.to(R).min(dim=1, keepdim=True)[0] / 2.0
    scale = scale.expand(-1, 2)
    c0 = image_size_wh / 2.0

    principal_point = -p0_pytorch3d * scale + c0
    focal_length = focal_pytorch3d * scale

    camera_matrix = torch.zeros_like(R)
    camera_matrix[:, :2, 2] = principal_point
    camera_matrix[:, 2, 2] = 1.0
    camera_matrix[:, 0, 0] = focal_length[:, 0]
    camera_matrix[:, 1, 1] = focal_length[:, 1]
    return R[0], tvec[0], camera_matrix[0]


def _load_16big_png_depth(depth_png):
    with Image.open(depth_png) as depth_pil:
        depth = np.array(depth_pil, dtype=np.uint16)
        depth = np.frombuffer(depth, dtype=np.float16).astype(np.float32)
        depth[np.isnan(depth)] = 0
        depth[np.isinf(depth)] = 0
        depth = depth.reshape((depth_pil.size[1], depth_pil.size[0]))

    return depth


def convert_ndc_to_pinhole(focal_length, principal_point, image_size):
    focal_length = np.array(focal_length)
    principal_point = np.array(principal_point)
    image_size_wh = np.array([image_size[1], image_size[0]])
    half_image_size = image_size_wh / 2
    rescale = half_image_size.min()
    principal_point_px = half_image_size - principal_point * rescale
    focal_length_px = focal_length * rescale
    fx, fy = focal_length_px[0], focal_length_px[1]
    cx, cy = principal_point_px[0], principal_point_px[1]
    K = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=np.float32)
    return K


def get_pose_and_K(filename, annotation):

    data = annotation['viewpoint']
    focal_length = data['focal_length']
    principal_points = data['principal_point']
    image_size = annotation['image']['size']

    R, tvec, camera_intrinsics = opencv_from_cameras_projection(
        np.array(np.array(data['R'])),
        np.array(np.array(data['T'])),
        np.array(focal_length),
        np.array(principal_points),
        np.array(image_size),
    )

    camera_pose = np.vstack([
        np.hstack((R, np.expand_dims(tvec, axis=1))),
        np.array([0.0, 0.0, 0.0, 1.0])
    ]).astype(np.float32)
    camera_pose = invert_extrinsics(camera_pose)

    return camera_pose, camera_intrinsics.numpy()


def prepare_annotations(seq, frame_annotations):
    annotations = {}
    for data in frame_annotations:
        frame_number = data['frame_number']
        sequence_name = data['sequence_name'] 
        if sequence_name not in annotations:
            annotations[sequence_name] = {}
        annotations[sequence_name][frame_number] = data
    sequence_name = seq.split('/')[-1]
    return annotations[sequence_name]

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

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


def get_depth(filename):
    return _load_16big_png_depth(filename)


def get_mask(filename):
    return np.array(read_image(filename, '1'))
    
#######################################################

def get_sequences(args):
    seqs = crawl(args.src, 'images')
    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','mask_rgb','mask_depth']

    ### Get filenames
    filename_rgbs = sorted(glob(f'{seq}/images/*'))
    filename_depths = sorted(glob(f'{seq}/depths/*'))
    filename_mask_depths = sorted(glob(f'{seq}/depth_masks/*'))
    filename_mask_rgbs = sorted(glob(f'{seq}/masks/*'))

    ### Prepare frame annotations
    frame_annotations = '/'.join(seq.split('/')[:-1]) + '/frame_annotations.jgz'
    frame_annotations = json.load(gzip.open(frame_annotations, mode="r"))
    annotations = prepare_annotations(seq, frame_annotations)

    ### Object name
    object_name = dst.split('/')[-2]

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

        ######## ANNOTATIONS
        for i, annotation in enumerate(annotations.values()):
            frame = frame_name(i)

            filename_rgb = f'{args.src}/{annotation["image"]["path"]}'
            filename_depth = f'{args.src}/{annotation["depth"]["path"]}'

            filename_mask_rgb = filename_rgb.replace('/images/', '/masks/').replace('.jpg', '.png')
            filename_mask_depth = filename_rgb.replace('/images/', '/depth_masks/').replace('.jpg', '.png')

            ######## RGB
            rgb = get_rgb(filename_rgb)
            dense['rgb'][frame] = rgb

            ######## LOWDIM RGB          
            prepare_lowdim(lowdim, dst, cam, frame)

            ######## MASK RGB
            mask_rgb = get_mask(filename_mask_rgb)
            dense['mask_rgb'][frame] = mask_rgb

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

            ######## MASK DEPTH
            mask_depth = get_mask(filename_mask_depth)
            dense['mask_depth'][frame] = mask_depth

            ######## INTRINSICS + EXTRINSICS
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            extrinsics, intrinsics = get_pose_and_K(filename_rgb, annotation)
            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='CO3Dv2',
            tags=['real','static','inward','object'],
            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=dict(object=object_name),
    )
    write_json(filename, seq_metadata)

    return dst

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

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

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