# Copyright 2026 Toyota Research Institute.  All rights reserved.

'''
Convert RAW Kubric-5D render output into the AnyData unified `videos` format.

Per scene `scn{idx:05d}/` the raw layout (from the Kubric-5D generation code) is:
    scn00000/
      frames_p0_v{v}/rgba_{t:05d}.png      # 16 views v=0..15, 60 frames t=0..59
      frames_p0_v{v}/depth_{t:05d}.tiff     # ray (point) distance, float meters
      scn00000_p0_v{v}.json                 # per-view camera (K, positions, quaternions) + scene meta

Output (one dir per scene): rgb/cam{vv}.mp4, depth/cam{vv}.zarr,
lowdim/cam{vv}.npz (intrinsics, extrinsics), metadata.json.

Camera and depth math are ported verbatim from the gcd Kubric loader
(dataproc/kubric4d.py:get_kubric_camera_matrices_numpy + load_scene scaling, and
dataproc/geometry.py:correct_depth_ball_plane_numpy) so the output reproduces the
existing unified Kubric-5D dataset. Structurally mirrors converters/processed.py.
'''

import os
import json
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,
)

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

NUM_VIEWS = 16                                          # cam{v} == raw frames_p0_v{v}
CAMERA_NAMES = [f'cam{v:02d}' for v in range(NUM_VIEWS)]   # cam00-03 high-45deg, cam04-15 low-5deg

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

def get_sequences(args):
    seqs = crawl(args.src, 'frames_p0_v0')              # marker dir present in every scene
    return sorted(os.path.dirname(seq) for seq in seqs)


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

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

def quat_to_rotmat(quat):
    '''
    Quaternion (Kubric / pyquaternion order w,x,y,z) to rotation matrix.
    Matches pyquaternion.Quaternion(quat).rotation_matrix used in gcd kubric4d.py.
    :param quat: (4,) array of float.
    :return R: (3, 3) array of float.
    '''
    q = np.asarray(quat, dtype=np.float64)
    q = q / (np.linalg.norm(q) + 1e-12)
    w, x, y, z = q
    return np.array([
        [1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - z * w),       2.0 * (x * z + y * w)],
        [2.0 * (x * y + z * w),       1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - x * w)],
        [2.0 * (x * z - y * w),       2.0 * (y * z + x * w),       1.0 - 2.0 * (x * x + y * y)],
    ], dtype=np.float64)


def camera_matrices(meta):
    '''
    Ported from gcd kubric4d.py:get_kubric_camera_matrices_numpy + load_scene intrinsic scaling.
    :param meta (dict): parsed per-view JSON.
    :return intrinsics: (T, 3, 3) float32, pixel-space pinhole.
    :return extrinsics: (T, 4, 4) float32, cam2world (columns right-down-forward).
    '''
    T = meta['scene']['num_frames']
    (W, H) = meta['scene']['resolution']                # raw JSON stores (width, height)

    # Normalized (sensor 1x1) intrinsics -> pixel space.
    K_norm = np.abs(np.array(meta['camera']['K'], dtype=np.float64))
    intrinsics = np.tile(K_norm[None], (T, 1, 1))
    intrinsics[:, 0, :] *= W
    intrinsics[:, 1, :] *= H

    extrinsics = []
    for t in range(T):
        E = np.eye(4, dtype=np.float64)
        E[0:3, 0:3] = quat_to_rotmat(meta['camera']['quaternions'][t])
        E[0:3, 3] = np.array(meta['camera']['positions'][t], dtype=np.float64)
        # Flip Y and Z columns: Blender (right-up-back) -> repo (right-down-forward).
        E[0:3, 1] *= -1.0
        E[0:3, 2] *= -1.0
        extrinsics.append(E)

    return intrinsics.astype(np.float32), np.stack(extrinsics, axis=0).astype(np.float32)


def correct_depth_ball_plane(depth_ball, intrinsics):
    '''
    Verbatim port of gcd geometry.py:correct_depth_ball_plane_numpy.
    Kubric stores ray (point-to-camera) distance; convert to planar z-depth.
    :param depth_ball: (H, W) array of float.
    :param intrinsics: (3, 3) array of float (pixel space).
    :return depth_plane: (H, W) array of float32.
    '''
    (H, W) = depth_ball.shape[-2:]
    fov_x = 2.0 * np.arctan(W / (2.0 * np.abs(intrinsics[0, 0])))
    fov_y = 2.0 * np.arctan(H / (2.0 * np.abs(intrinsics[1, 1])))
    angles_x = np.linspace(-fov_x / 2.0, fov_x / 2.0, W)
    angles_y = np.linspace(-fov_y / 2.0, fov_y / 2.0, H)
    mismatch_x = np.tan(angles_x)
    mismatch_y = np.tan(angles_y)
    correction = np.sqrt(mismatch_x[None, :] ** 2 + mismatch_y[:, None] ** 2 + 1.0)
    return (depth_ball / correction).astype(np.float32)

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

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

    scn = os.path.basename(seq)                         # 'scn00000'

    num_frames = {cam: dict() for cam in CAMERA_NAMES}
    resolution = {cam: dict() for cam in CAMERA_NAMES}
    labels, lowdim = [], {}

    ############ LOOP OVER CAMERAS
    for v, cam in enumerate(CAMERA_NAMES):
        dense = {'rgb': dict(), 'depth': dict()}

        frame_dir = f'{seq}/frames_p0_v{v}'
        view_json = f'{seq}/{scn}_p0_v{v}.json'
        filename_rgbs = sorted(glob(f'{frame_dir}/rgba_*.png'))
        filename_depths = sorted(glob(f'{frame_dir}/depth_*.tiff'))
        # filename_segm = sorted(glob(f'{frame_dir}/segmentation_*.png'))   # OPTIONAL, see note below

        if len(filename_rgbs) == 0:
            continue

        with open(view_json, 'r') as f:
            meta = json.load(f)
        intrinsics, extrinsics = camera_matrices(meta)              # (T,3,3), (T,4,4) cam2world

        ######## RGB + per-frame LOWDIM (intrinsics, extrinsics)
        for t, filename_rgb in enumerate(filename_rgbs):
            frame = frame_name(t)
            rgb = np.asarray(read_image(filename_rgb))[..., :3]     # (H,W,3) uint8, drop alpha
            dense['rgb'][frame] = rgb

            prepare_lowdim(lowdim, dst, cam, frame)                 # seeds camera, timestep
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['intrinsics'] = intrinsics[t]
            lowdim[filename_lowdim]['extrinsics'] = extrinsics[t]

        ######## DEPTH (ray meters -> planar meters; Kubric-5D is dome-bounded, range ~[5.5, 55], no sentinel)
        for t, filename_depth in enumerate(filename_depths):
            frame = frame_name(t)
            depth_ball = np.asarray(Image.open(filename_depth), dtype=np.float32)   # (H,W)
            dense['depth'][frame] = correct_depth_ball_plane(depth_ball, intrinsics[t])

        ######## OPTIONAL: Kubric instance segmentation (8-bit zarr; not in the released 4-label set).
        # AnyData supports 'semantic' (BIT_COMPRESS_LABELS), but normals / object_coordinates /
        # pointmap have no first-class write support and are derivable from depth+intrinsics+extrinsics,
        # so they are omitted. To enable segmentation, add 'semantic' frames here, e.g.:
        # for t, fn in enumerate(filename_segm):
        #     dense.setdefault('semantic', {})[frame_name(t)] = np.asarray(Image.open(fn))

        write_labels(dst, cam, args.storage, dense, labels, resolution, num_frames)

    ######## WRITE STACKED LOWDIM (one npz per camera: intrinsics (T,3,3), extrinsics (T,4,4))
    write_lowdim(args, dst, labels, num_frames, lowdim)

    ############ METADATA
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='Kubric5D',
            tags=['dynamic', 'indoors', 'inward', 'sim'],
            raw_id=scn,
        ),
        labels=labels,
        cameras=CAMERA_NAMES,
        resolution=resolution,
        num_frames=num_frames,
        framerate=24,
        rgb=dict(extension='jpg'),                                  # forced to mp4 in videos mode
        intrinsics=dict(model='pinhole'),
        extrinsics=dict(transform='cam2world', metric=False),
        depth=dict(extension='npz', metric=False, sparse=False),    # forced to zarr in videos mode
        semantic=None,
        action=None,
        language=None,
        specific=None,
    )
    write_json(f'{dst}/metadata.json', seq_metadata)

    return dst

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

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

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