# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import cv2
import numpy as np
import pyexr
import torch

from PIL import Image
from glob import glob
import json

from anydata.geometry.camera import Camera
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

os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"

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

def load_exr(filename, resize=(256, 256)):
    f = pyexr.open(filename)
    img = f.get("default")

    alpha = img[..., 3:4]
    img[..., :3] *= alpha
    img[..., :3] += (1 - alpha)
    img = np.clip(img, 0.0, 1.0)

    if resize is not None:
        img = resize_mip(img, resize, interpolation=cv2.INTER_AREA)

    img = torch.from_numpy(img)[..., :3]
    img = linear_to_srgb(img)

    return img


def linear_to_srgb(img):
    limit = 0.0031308
    img = torch.where(img > limit, 1.055 * (img ** (1.0 / 2.4)) - 0.055, 12.92 * img)
    img[img > 1] = 1
    return img


def resize_mip(img, resize=[256, 256], interpolation=cv2.INTER_LINEAR):
    shape = (resize[1], resize[0])
    img = cv2.resize(img, dsize=shape, interpolation=interpolation)
    return img


def get_rgb(filename):
    rgb = load_exr(filename, resize=(1600, 1600))
    rgb = Image.fromarray((rgb.numpy() * 255).astype(np.uint8))
    return np.array(rgb)


def get_intrinsics(data, scale=1.0):
    intrinsics = np.eye(3)
    intrinsics_raw = data["camera_data"]["intrinsics"]
    intrinsics[0, 0] = intrinsics_raw['fx'] * scale
    intrinsics[1, 1] = intrinsics_raw['fy'] * scale
    intrinsics[0, 2] = intrinsics_raw['cx'] * scale
    intrinsics[1, 2] = intrinsics_raw['cy'] * scale
    return intrinsics


def get_extrinsics(data):
    convert_matrix = torch.tensor(
        [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
        dtype=torch.float32).numpy()
    c2w = np.array(data["camera_data"]["cam2world"]).T
    pose = convert_matrix @ c2w @ convert_matrix
    return pose


def get_depth(filename, intrinsics):
    depth = cv2.imread(filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)[:, :, 0]
    depth[depth < 0] = 0
    depth = torch.tensor(depth).unsqueeze(0).unsqueeze(0)
    cam = Camera(K=torch.tensor(intrinsics).unsqueeze(0).float(), hw=depth.shape[-2:])
    points = cam.reconstruct_depth_map_rays(depth, to_world=False)
    return points[0, 2, :, :].numpy()

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

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


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

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

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']

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

        filenames = sorted(glob(f'{seq}/*.exr'))
        filename_rgbs = [f for f in filenames if '.depth.' not in f and '.seg.' not in f]
        filename_depths = [f for f in filenames if '.depth.' in f]

        filename_datas = sorted(glob(f'{seq}/*.json'))
        filename_datas = [f for f in filename_datas if 'nerf_' not in f]

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

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

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

        ######## DATA FILENAMES
        all_intrinsics = []
        for i, filename_data in enumerate(filename_datas):
            frame = frame_name(i)

            with open(filename_data, 'r') as f:
                data = json.loads(f.read())

            ######## INTRINSICS + EXTRINSICS
            extrinsics = get_extrinsics(data)
            intrinsics = get_intrinsics(data)
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['extrinsics'] = extrinsics
            lowdim[filename_lowdim]['intrinsics'] = intrinsics
            all_intrinsics.append(intrinsics)

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

            ######## DEPTH
            depth = get_depth(filename_depth, all_intrinsics[i])
            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='RTMV',
            tags=['sim','static','objects','image'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=0,
        rgb=dict(extension='jpg'),
        intrinsics=dict(model='pinhole'),
        extrinsics=dict(transform='cam2world',metric=False),
        depth=dict(extension='npz',metric=False,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)

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