# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import h5py
import numpy as np
import pandas as pd 

from PIL import Image
from glob import glob
from pylab import imsave, clip

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


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

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

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

    dsts = []
    for opt in range(10):

        cameras = ['0']
        num_frames = {cam: dict() for cam in cameras}
        resolution = {cam: dict() for cam in cameras}
        labels, lowdim = [], {}
        dense_labels = ['rgb','depth']

        cam = cameras[0]
        dense = {label: dict() for label in dense_labels}
        opt = '%02d' % opt

        ### Get Filenames
        filename_rgbs = sorted(glob(f'{seq}/scene_cam_{opt}_final_hdf5/*.color.hdf5'))
        filename_depths = sorted(glob(f'{seq}/scene_cam_{opt}_geometry_hdf5/*.depth_meters.hdf5'))
        filename_positions = f'{os.path.dirname(seq)}/_detail/cam_{opt}/camera_keyframe_positions.hdf5'
        filename_orientations = f'{os.path.dirname(seq)}/_detail/cam_{opt}/camera_keyframe_orientations.hdf5'
        filename_metadata = f'{os.path.dirname(seq)}/_detail/metadata_scene.csv'

        if len(filename_rgbs) == 0:
            continue

        dst = f'{dst_all}/%02d' % len(dsts)
        dsts.append(dst)

        gamma                             = 1.0 / 2.2
        inv_gamma                         = 1.0 / gamma
        percentile                        = 90
        brightness_nth_percentile_desired = 0.8
        eps                               = 0.0001

        scene_name = seq.split('/')[-2]
        metadata_path = f'{args.src}/metadata_camera_parameters.csv'
        df_camera_parameters = pd.read_csv(metadata_path, index_col="scene_name")

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

            with h5py.File(filename_rgb, "r") as f: 
                rgb = f["dataset"][:].astype(np.float32)

            brightness = 0.3 * rgb[:,:,0] + 0.59 * rgb[:,:,1] + 0.11 * rgb[:,:,2]
            brightness_nth_percentile_current = np.percentile(brightness, percentile) + 1e-6
            scale = np.power(brightness_nth_percentile_desired, inv_gamma) / brightness_nth_percentile_current
            rgb = np.power(np.maximum(scale * rgb, 0), gamma)
            dense['rgb'][frame] = (clip(rgb, 0, 1) * 255).astype(np.uint8)

            # if 'rgb' not in labels: labels.append('rgb')
            # filename_rgb_out = f'{dst}/rgb/{cam}/{frame}.jpg'
            # create_folder(filename_rgb_out)
            # imsave(filename_rgb_out, clip(rgb, 0, 1))

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

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

            ### DEPTH
            with h5py.File(filename_depth, "r") as f: depth = f["dataset"][:].astype(np.float32)
            filename_depth_out = f'{dst}/depth/{cam}/{frame}.npz'

            intWidth, intHeight = 1024, 768
            fltFocal = 886.81

            import numpy
            npyImageplaneX = numpy.linspace((-0.5 * intWidth) + 0.5, (0.5 * intWidth) - 0.5, intWidth).reshape(1, intWidth).repeat(intHeight, 0).astype(numpy.float32)[:, :, None]
            npyImageplaneY = numpy.linspace((-0.5 * intHeight) + 0.5, (0.5 * intHeight) - 0.5, intHeight).reshape(intHeight, 1).repeat(intWidth, 1).astype(numpy.float32)[:, :, None]
            npyImageplaneZ = numpy.full([intHeight, intWidth, 1], fltFocal, numpy.float32)
            npyImageplane = numpy.concatenate([npyImageplaneX, npyImageplaneY, npyImageplaneZ], 2)
            depth = depth / numpy.linalg.norm(npyImageplane, 2, 2) * fltFocal
            depth[np.isnan(depth)] = 0

            dense['depth'][frame] = depth

            with h5py.File(filename_positions,    "r") as f: camera_positions    = f["dataset"][i]
            with h5py.File(filename_orientations, "r") as f: camera_orientations = f["dataset"][i]

            # %% Get scale factor from asset coordinates to meters
            scene_metadata = pd.read_csv(filename_metadata)
            meters_per_asset_unit = scene_metadata[
                scene_metadata.parameter_name == "meters_per_asset_unit"
            ]
            assert len(meters_per_asset_unit) == 1  # Should not be multiply defined
            meters_per_asset_unit = meters_per_asset_unit.parameter_value[0]
            scale_factor = float(meters_per_asset_unit)

            # %% Get camera pose and intrinsics
            camera_position_world = camera_positions
            R_world_from_cam = camera_orientations
            t_world_from_cam = np.array(camera_position_world).T

            M_world_from_cam = np.zeros((4, 4))
            M_world_from_cam[:3, :3] = R_world_from_cam
            M_world_from_cam[:3, 3] = t_world_from_cam * scale_factor
            M_world_from_cam[3, 3] = 1.0

            world_T_cam = M_world_from_cam.astype(np.float32)

            gl_to_cv = np.array([[1, -1, -1, 1], [-1, 1, 1, -1], [-1, 1, 1, -1], [1, 1, 1, 1]])

            rotx = np.array([
                [1, 0, 0],[0,1,0],[0,0,1],
            ])

            world_T_cam *= gl_to_cv

            rot_mat = world_T_cam[:3, :3]
            trans = world_T_cam[:3, 3]

            # rot_mat = rotx(-np.pi / 2) @ rot_mat
            # trans = rotx(-np.pi / 2) @ trans

            world_T_cam[:3, :3] = rot_mat
            world_T_cam[:3, 3] = trans

            cam_T_world = np.linalg.inv(world_T_cam)

            df_ = df_camera_parameters.loc[scene_name]

            width_pixels = int(df_["settings_output_img_width"])
            height_pixels = int(df_["settings_output_img_height"])

            M_proj = np.array(
                [
                    [df_["M_proj_00"], df_["M_proj_01"], df_["M_proj_02"], df_["M_proj_03"]],
                    [df_["M_proj_10"], df_["M_proj_11"], df_["M_proj_12"], df_["M_proj_13"]],
                    [df_["M_proj_20"], df_["M_proj_21"], df_["M_proj_22"], df_["M_proj_23"]],
                    [df_["M_proj_30"], df_["M_proj_31"], df_["M_proj_32"], df_["M_proj_33"]],
                ]
            )

            # matrix to map to integer screen coordinates from normalized device coordinates
            M_screen_from_ndc = np.matrix(
                [
                    [0.5 * (width_pixels - 1), 0, 0, 0.5 * (width_pixels - 1)],
                    [0, -0.5 * (height_pixels - 1), 0, 0.5 * (height_pixels - 1)],
                    [0, 0, 0.5, 0.5],
                    [0, 0, 0, 1.0],
                ]
            )

            # Extract fx, fy, cx and cy, and build a corresponding matrix
            screen_from_cam = M_screen_from_ndc @ M_proj
            fx = abs(screen_from_cam[0, 0])
            fy = abs(screen_from_cam[1, 1])
            cx = abs(screen_from_cam[0, 2])
            cy = abs(screen_from_cam[1, 2])

            K = np.eye(3, dtype=np.float32)
            K[0, 0] = float(fx)
            K[1, 1] = float(fy)
            K[0, 2] = float(cx)
            K[1, 2] = float(cy)

            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['extrinsics'] = invert_extrinsics(cam_T_world)
            lowdim[filename_lowdim]['intrinsics'] = K

        ######## 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='Hypersim',
                tags=['sim','static','indoor'],
                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 dsts

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

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

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