import os
import numpy as np
import torch
import torch.utils.data as data
import torch.nn.functional as F
import logging
import os
import re
import copy
import math
import random
from pathlib import Path
from glob import glob
import os.path as osp

from core.utils import frame_utils

class CREStereo():
    def __init__(self, aug_params=None, root='/data/CREStereo'):
        self.root = root
        assert os.path.exists(root)

        disp_list = self.selector('_left.disp.png')
        image1_list = self.selector('_left.jpg')
        image2_list = self.selector('_right.jpg')

        assert len(image1_list) == len(image2_list) == len(disp_list) > 0
        for img1, img2, disp in zip(image1_list, image2_list, disp_list):
            # if random.randint(1, 20000) != 1:
            #     continue
            self.image_list += [[img1, img2]]
            self.disparity_list += [disp]

    def selector(self, suffix):
        files = list(glob(os.path.join(self.root, f"hole/*{suffix}")))
        files += list(glob(os.path.join(self.root, f"shapenet/*{suffix}")))
        files += list(glob(os.path.join(self.root, f"tree/*{suffix}")))
        files += list(glob(os.path.join(self.root, f"reflective/*{suffix}")))
        return sorted(files)

    def disparity_distribution(self):
        disp_lists = self.selector('_left.disp.png')
        disparities = []
        for filename in disp_lists:
            disp_gt = frame_utils.readDispCREStereo(filename)
            print(disp_gt.shape)


CREStereo.disparity_distribution()