import cv2
import numpy as np
from models.SpaTrackV2.datasets.dataset_util import imread_cv2
from models.SpaTrackV2.datasets.co3d_v2 import Co3d
import os.path as osp

class WildRGBD(Co3d):
    def __init__(self, mask_bg=True, *args, ROOT, **kwargs):
        super().__init__(mask_bg, *args, ROOT=ROOT, **kwargs)
        self.dataset_label = 'WildRGBD'

    def _get_metadatapath(self, obj, instance, view_idx):
        return osp.join(self.ROOT, obj, instance, 'metadata', f'{view_idx:0>5d}.npz')

    def _get_impath(self, obj, instance, view_idx):
        return osp.join(self.ROOT, obj, instance, 'rgb', f'{view_idx:0>5d}.jpg')

    def _get_depthpath(self, obj, instance, view_idx):
        return osp.join(self.ROOT, obj, instance, 'depth', f'{view_idx:0>5d}.png')

    def _get_maskpath(self, obj, instance, view_idx):
        return osp.join(self.ROOT, obj, instance, 'masks', f'{view_idx:0>5d}.png')

    def _read_depthmap(self, depthpath, input_metadata):
        # We store depths in the depth scale of 1000.
        # That is, when we load depth image and divide by 1000, we could get depth in meters.
        depthmap = imread_cv2(depthpath, cv2.IMREAD_UNCHANGED)
        depthmap = depthmap.astype(np.float32) / 1000.0
        return depthmap

if __name__ == "__main__":
    from models.videocam.datasets.base_sfm_dataset import view_name
    from functools import partial
    import os
    # from dust3r.viz import SceneViz, auto_cam_size
    # from dust3r.utils.image import rgb
    DATA_DIR = "/nas6/datasets/wildrgbd_processed"
        
    dataset = WildRGBD(split='train', ROOT=DATA_DIR, mask_bg=True, 
                       resolution=384, aug_crop=16, num_views=16)
    rng = np.random.default_rng(seed=0)
    data_ret = dataset._get_views(1500, (518,518), rng)
    # prefetch_data = partial(dataset.prefetch_data, cpu_num_avg=cpu_num_per, bakup_dir=bakup_dir)
    # scene_idx_list = [i for i in range(len(dataset))]
    # if sys.argv[3] == "true":
    #     with concurrent.futures.ProcessPoolExecutor(max_workers=cpu_num_total//20) as executor:
    #         pool_buffer = executor.map(prefetch_data, scene_idx_list) #, total=len(dataset))
    # elif sys.argv[3] == "false":   
    #     for idx in tqdm.tqdm(range(len(dataset))):
    #         prefetch_data(idx)  
    from models.videocam.datasets.vis3d_check import vis4d
    os.system(f"rm -rf /home/xyx/home/codes/SpaTrackerV2/vis_results/test")
    vis4d(data_ret["rgbs"], data_ret["depths"],
           data_ret["traj_mat"], data_ret["intrs"], workspace="/home/xyx/home/codes/SpaTrackerV2/vis_results/test")
    import pdb; pdb.set_trace()