# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Dataloader for preprocessed StaticThings3D
# dataset at https://github.com/lmb-freiburg/robustmvd/
# See datasets_preprocess/preprocess_staticthings3d.py
# --------------------------------------------------------
import os.path as osp
import numpy as np

from dust3r.datasets.base.base_stereo_view_dataset import BaseStereoViewDataset
from dust3r.utils.image import imread_cv2


class StaticThings3D(BaseStereoViewDataset):
    """Dataset of indoor scenes, 5 images each time"""

    def __init__(self, ROOT, *args, mask_bg="rand", **kwargs):
        self.ROOT = ROOT
        super().__init__(*args, **kwargs)

        assert mask_bg in (True, False, "rand")
        self.mask_bg = mask_bg

        # loading all pairs
        assert self.split is None
        self.pairs = np.load(osp.join(ROOT, "staticthings_pairs.npy"))

    def __len__(self):
        return len(self.pairs)

    def get_stats(self):
        return f"{len(self)} pairs"

    def _get_views(self, pair_idx, resolution, rng):
        scene, seq, cam1, im1, cam2, im2 = self.pairs[pair_idx]
        seq_path = osp.join("TRAIN", scene.decode("ascii"), f"{seq:04d}")

        views = []

        mask_bg = (self.mask_bg == True) or (self.mask_bg == "rand" and rng.choice(2))

        CAM = {b"l": "left", b"r": "right"}
        for cam, idx in [(CAM[cam1], im1), (CAM[cam2], im2)]:
            num = f"{idx:04n}"
            img = num + "_clean.jpg" if rng.choice(2) else num + "_final.jpg"
            image = imread_cv2(osp.join(self.ROOT, seq_path, cam, img))
            depthmap = imread_cv2(osp.join(self.ROOT, seq_path, cam, num + ".exr"))
            camera_params = np.load(osp.join(self.ROOT, seq_path, cam, num + ".npz"))

            intrinsics = camera_params["intrinsics"]
            camera_pose = camera_params["cam2world"]

            if mask_bg:
                depthmap[depthmap > 200] = 0

            image, depthmap, intrinsics = self._crop_resize_if_necessary(
                image, depthmap, intrinsics, resolution, rng, info=(seq_path, cam, img)
            )

            views.append(
                dict(
                    img=image,
                    depthmap=depthmap,
                    camera_pose=camera_pose,  # cam2world
                    camera_intrinsics=intrinsics,
                    dataset="StaticThings3D",
                    label=seq_path,
                    instance=cam + "_" + img,
                )
            )

        return views


if __name__ == "__main__":
    from dust3r.datasets.base.base_stereo_view_dataset import view_name
    from dust3r.viz import SceneViz, auto_cam_size
    from dust3r.utils.image import rgb

    dataset = StaticThings3D(ROOT="data/staticthings3d_processed", resolution=224, aug_crop=16)

    for idx in np.random.permutation(len(dataset)):
        views = dataset[idx]
        assert len(views) == 2
        print(idx, view_name(views[0]), view_name(views[1]))
        viz = SceneViz()
        poses = [views[view_idx]["camera_pose"] for view_idx in [0, 1]]
        cam_size = max(auto_cam_size(poses), 0.001)
        for view_idx in [0, 1]:
            pts3d = views[view_idx]["pts3d"]
            valid_mask = views[view_idx]["valid_mask"]
            colors = rgb(views[view_idx]["img"])
            viz.add_pointcloud(pts3d, colors, valid_mask)
            viz.add_camera(
                pose_c2w=views[view_idx]["camera_pose"],
                focal=views[view_idx]["camera_intrinsics"][0, 0],
                color=(idx * 255, (1 - idx) * 255, 0),
                image=colors,
                cam_size=cam_size,
            )
        viz.show()
