from __future__ import absolute_import, division, print_function

import os

import matplotlib.pyplot as plt

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import random
import numpy as np
from PIL import Image  # using pillow-simd for increased speed
import random
import numpy as np
from PIL import Image  # using pillow-simd for increased speed
from skimage.filters import gaussian, sobel
from scipy.interpolate import griddata
import cv2
from utils import transfer_color
import torch
import torch.utils.data as data
from torchvision import transforms



def pil_loader(path):
    # print(len(path))
    with open(path, 'rb') as f:
        with Image.open(f) as img:
            return img.convert('RGB')


def read_pfm(file):
    with open(file, 'rb') as fh:
        fh.readline()
        width, height = str(fh.readline().rstrip())[2:-1].split()
        fh.readline()
        disp = np.fromfile(fh, '<f')
        return np.flipud(disp.reshape(int(height), int(width)))


class BaseDataset:
    def __init__(self,
                 data_path,
                 feed_height=1024,
                 feed_width=768,
                 max_disparity=20,
                 is_train=True,
                 disable_synthetic_augmentation=True,
                 disable_sharpening=True,
                 monodepth_model='midas',
                 disable_background=False,
                 has_gt=True,
                 disable_normalisation=False,
                 keep_aspect_ratio=True):
        
        itxt_path = '/home/notebook/data/group/Ruby/stereo-from-mono-master/img.txt'
        dtxt_path = '/home/notebook/data/group/Ruby/stereo-from-mono-master/depth.txt'
        all_img = []
        all_depth = []
        with open( itxt_path , 'r' ) as fi:
            for li in fi.readlines():
                li = li.strip( '\n' )
                all_img.append(li)

            # print('all_image_len',len(all_img[0]))
            # print('all_image_len',all_img[0])
        with open( dtxt_path , 'r' ) as fd:
            for ld in fd.readlines():
                ld = ld.strip( '\n' )
                all_depth.append(ld)

            # print('all_depth',len(all_depth[0]))
                  
        self.data_path = data_path
        # self.depth_path=self.data_path+"/"
        # self.img_path=self.data_path+"/"
        self.depth_path= all_depth
        self.img_path= all_img
        # self.filenames = filenames
        self.feed_height = feed_height
        self.feed_width = feed_width
        self.is_train = is_train
        self.has_gt = has_gt
        self.disable_normalisation = disable_normalisation
        self.keep_aspect_ratio = keep_aspect_ratio

        self.loader = pil_loader
        self.read_pfm = read_pfm
        self.to_tensor = transforms.ToTensor()
        # self.filenames=[file for file in os.listdir(self.img_path)  if(".png"  in file)]
        # self.filenames=[file for file in os.listdir(self.depth_path)  if(".png"  in file)]
        self.max_disparity = max_disparity
        self.disable_synthetic_augmentation = disable_synthetic_augmentation
        self.disable_sharpening = disable_sharpening
        self.monodepth_model = monodepth_model
        self.disable_background = disable_background

        # do image generation for a wider image so we can crop off missing pixels
        self.process_width = self.feed_width + self.max_disparity

        self.xs, self.ys = np.meshgrid(np.arange(self.process_width), np.arange(self.feed_height))

        # We need to specify augmentations differently in newer versions of torchvision.
        # We first try the newer tuple version; if this fails we fall back to scalars
        try:
            self.stereo_brightness = (0.8, 1.2)
            self.stereo_contrast = (0.8, 1.2)
            self.stereo_saturation = (0.8, 1.2)
            self.stereo_hue = (-0.01, 0.01)
            transforms.ColorJitter.get_params(
                self.stereo_brightness, self.stereo_contrast, self.stereo_saturation,
                self.stereo_hue)
        except TypeError:
            self.stereo_brightness = 0.2
            self.stereo_contrast = 0.2
            self.stereo_saturation = 0.2
            self.stereo_hue = 0.01
        try:
            self.brightness = (0.8, 1.2)
            self.contrast = (0.8, 1.2)
            self.saturation = (0.8, 1.2)
            self.hue = (-0.1, 0.1)
            transforms.ColorJitter.get_params(
                self.brightness, self.contrast, self.saturation, self.hue)
        except TypeError:
            self.brightness = 0.2
            self.contrast = 0.2
            self.saturation = 0.2
            self.hue = 0.1

        self.silly_svsm = False

    def __len__(self):
        # return len(self.filenames)
        return len(self.img_path)


    def load_images(self, idx, do_flip=False):
        """ Load an image to use as left and a random background image to fill in occlusion holes"""

        # frame = self.filenames[idx]
        # image = self.loader(os.path.join(self.img_path, frame))
        # print('img_path',self.img_path)
        ###这边是单张数据读取
        print('idx',idx)
        print('img',self.img_path[idx])
        image = self.loader(self.img_path[idx])

        if do_flip:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
        # bg_path=r"/home/notebook/data/group/Ruby/stereo-from-mono-master/flower/img"
        # frame = random.choice(os.listdir(bg_path))
        frame = random.choice(self.img_path)
        # background = self.loader(os.path.join(bg_path, frame))
        background = self.loader(frame)
        return image, background

    def load_disparity(self, idx, do_flip=False):
        # frame = self.filenames[idx]
        # print(self.depth_path+frame.replace(".png","_depth.png"))
        # depth=self.loader(self.depth_path+frame.replace(".png","_depth.png")).convert("L")
        depth=self.loader(self.depth_path[idx]).convert("L")
        disparity = np.squeeze(depth)
        # disparity=25-disparity

        if do_flip:
            disparity = disparity[:, ::-1]
        return disparity

    def preprocess(self, inputs):

        # do color augmentation
        if self.is_train and random.random() > 0.5:
            color_aug = transforms.ColorJitter.get_params(
                self.brightness, self.contrast, self.saturation, self.hue)

            for key in ['image', 'stereo_image']:
                inputs[key] = color_aug(inputs[key])

        # convert to tensors and standardise using ImageNet
        for key in ['image', 'stereo_image']:
            if self.disable_normalisation:
                inputs[key] = self.to_tensor(inputs[key])
            else:
                inputs[key] = (self.to_tensor(inputs[key]) - 0.45) / 0.225
        if self.has_gt:
            inputs['disparity'] = torch.from_numpy(inputs['disparity']).float()

        if inputs.get('mono_disparity') is not None:
            inputs['mono_disparity'] = torch.from_numpy(inputs['mono_disparity']).float()
    def process_disparity(self, disparity, max_disparity_range=(40, 196)):
        """ Depth predictions have arbitrary scale - need to convert to a pixel disparity"""

        disparity = disparity.copy()

        # make disparities positive
        min_disp = disparity.min()
        if min_disp < 0:
            disparity += np.abs(min_disp)

        # if random.random() < 0.01:
        #     # make max warped disparity bigger than network max -> will be clipped to max disparity,
        #     # but will mean network is robust to disparities which are too big
        #     max_disparity_range = (self.max_disparity * 1.05, self.max_disparity * 1.15)

        disparity /= disparity.max()  # now 0-1

        scaling_factor = (max_disparity_range[0] + random.random() *
                          (max_disparity_range[1] - max_disparity_range[0]))
        disparity *= scaling_factor

        if not self.disable_sharpening:
            # now find disparity gradients and set to nearest - stop flying pixels
            edges = sobel(disparity) > 3
            disparity[edges] = 0
            mask = disparity > 0

            try:
                disparity = griddata(np.stack([self.ys[mask].ravel(), self.xs[mask].ravel()], 1),
                                     disparity[mask].ravel(), np.stack([self.ys.ravel(),
                                                                        self.xs.ravel()], 1),
                                     method='nearest').reshape(self.feed_height, self.process_width)
            except (ValueError, IndexError) as e:
                pass  # just return disparity

        return disparity

    def prepare_sizes(self, inputs):

        height, width, _ = np.array(inputs['left_image']).shape

        if self.keep_aspect_ratio:
            if self.feed_height <= height and self.process_width <= width:
                # can simply crop the image
                target_height = height
                target_width = width

            else:
                # check the constraint
                current_ratio = height / width
                target_ratio = self.feed_height / self.process_width

                if current_ratio < target_ratio:
                    # height is the constraint
                    target_height = self.feed_height
                    target_width = int(self.feed_height / height * width)
                elif current_ratio > target_ratio:
                    # width is the constraint
                    target_height = int(self.process_width / width * height)
                    target_width = self.process_width
                else:
                    # ratio is the same - just resize
                    target_height = self.feed_height
                    target_width = self.process_width

        else:
            target_height = self.feed_height
            target_width = self.process_width

        inputs = self.resize_all(inputs, target_height, target_width)

        # now do cropping
        if target_height == self.feed_height and target_width == self.process_width:
            # we are already at the correct size - no cropping
            pass
        else:
            self.crop_all(inputs)

        return inputs

    def crop_all(self, inputs):

        # get crop parameters
        height, width, _ = np.array(inputs['left_image']).shape
        top = int(random.random() * (height - self.feed_height))
        left = int(random.random() * (width - self.process_width))
        right, bottom = left + self.process_width, top + self.feed_height

        for key in ['left_image', 'background']:
            inputs[key] = inputs[key].crop((left, top, right, bottom))
        inputs['loaded_disparity'] = inputs['loaded_disparity'][top:bottom, left:right]

        return inputs

    @staticmethod
    def resize_all(inputs, height, width):

        # images
        img_resizer = transforms.Resize(size=(height, width))
        for key in ['left_image', 'background']:
            inputs[key] = img_resizer(inputs[key])
        # disparity - needs rescaling
        disp = inputs['loaded_disparity']
        disp=np.array(disp).astype(np.float32)
        disp *= width / disp.shape[1]

        disp = cv2.resize(disp.astype(float), (width, height))  # ensure disp is float32 for cv2
        inputs['loaded_disparity'] = disp

        return inputs

    def get_occlusion_mask(self, shifted):

        mask_up = shifted > 0
        mask_down = shifted > 0

        shifted_up = np.ceil(shifted)
        shifted_down = np.floor(shifted)

        for col in range(self.process_width - 2):
            loc = shifted[:, col:col + 1]  # keepdims
            loc_up = np.ceil(loc)
            loc_down = np.floor(loc)

            _mask_down = ((shifted_down[:, col + 2:] != loc_down) * (
            (shifted_up[:, col + 2:] != loc_down))).min(-1)
            _mask_up = ((shifted_down[:, col + 2:] != loc_up) * (
            (shifted_up[:, col + 2:] != loc_up))).min(-1)

            mask_up[:, col] = mask_up[:, col] * _mask_up
            mask_down[:, col] = mask_down[:, col] * _mask_down

        mask = mask_up + mask_down
        return mask

    def project_image(self, image, disp_map, background_image):

        image = np.array(image)
        background_image = np.array(background_image)

        # set up for projection
        warped_image = np.zeros_like(image).astype(float)
        warped_image = np.stack([warped_image] * 2, 0)
        pix_locations = self.xs - disp_map

        # find where occlusions are, and remove from disparity map
        mask = self.get_occlusion_mask(pix_locations)
        masked_pix_locations = pix_locations * mask - self.process_width * (1 - mask)

        # do projection - linear interpolate up to 1 pixel away
        weights = np.ones((2, self.feed_height, self.process_width)) * 10000

        for col in range(self.process_width - 1, -1, -1):
            loc = masked_pix_locations[:, col]
            loc_up = np.ceil(loc).astype(int)
            loc_down = np.floor(loc).astype(int)
            weight_up = loc_up - loc
            weight_down = 1 - weight_up

            mask = loc_up >= 0
            mask[mask] = \
                weights[0, np.arange(self.feed_height)[mask], loc_up[mask]] > weight_up[mask]
            weights[0, np.arange(self.feed_height)[mask], loc_up[mask]] = \
                weight_up[mask]
            warped_image[0, np.arange(self.feed_height)[mask], loc_up[mask]] = \
                image[:, col][mask] / 255.

            mask = loc_down >= 0
            mask[mask] = \
                weights[1, np.arange(self.feed_height)[mask], loc_down[mask]] > weight_down[mask]
            weights[1, np.arange(self.feed_height)[mask], loc_down[mask]] = weight_down[mask]
            warped_image[1, np.arange(self.feed_height)[mask], loc_down[mask]] = \
                image[:, col][mask] / 255.

        weights /= weights.sum(0, keepdims=True) + 1e-7  # normalise
        weights = np.expand_dims(weights, -1)
        warped_image = warped_image[0] * weights[1] + warped_image[1] * weights[0]
        warped_image *= 255.

        # now fill occluded regions with random background
        if not self.disable_background:
            warped_image[warped_image.max(-1) == 0] = background_image[warped_image.max(-1) == 0]

        warped_image = warped_image.astype(np.uint8)

        return warped_image

    def augment_synthetic_image(self, image):

        if self.disable_synthetic_augmentation:

            return Image.fromarray(image.astype(np.uint8))

        # add some noise to stereo image
        noise = np.random.randn(self.feed_height, self.process_width, 3) / 50
        image = np.clip(image / 255 + noise, 0, 1) * 255

        # add blurring
        if random.random() > 0.5:
            image = gaussian(image,
                             sigma=random.random())
        image = np.clip(image, 0, 255)

        # color augmentation
        # stereo_aug = transforms.ColorJitter.get_params(
        #     self.stereo_brightness, self.stereo_contrast, self.stereo_saturation,
        #     self.stereo_hue)
        # print(image.shape)
        # image = stereo_aug(Image.fromarray(image.astype(np.uint8)))
        image=Image.fromarray(image.astype(np.uint8))
        return image

    def __getitem__(self, idx):
        inputs = {}
        do_flip = False
        if self.is_train and random.random() > 0.5:
            do_flip = True
        # load from disk
        left_image, background_image = self.load_images(idx, do_flip=do_flip)
        loaded_disparity = self.load_disparity(idx, do_flip=do_flip)

        inputs['left_image'] = left_image
        inputs['background'] = background_image
        inputs['loaded_disparity'] = loaded_disparity

        # resize and/or crop
        inputs = self.prepare_sizes(inputs)

        # match color in background image
        inputs['background'] = transfer_color(np.array(inputs['background']),
                                              np.array(inputs['left_image']))

        # convert scaleless disparity to pixel disparity
        inputs['disparity'] = \
            self.process_disparity(inputs['loaded_disparity'],
                                   max_disparity_range=(10, self.max_disparity))

        # now generate synthetic stereo image###
        projection_disparity = inputs['disparity']
        right_image = self.project_image(inputs['left_image'],
                                         projection_disparity, inputs['background'])

        # augmentation ###
        right_image = self.augment_synthetic_image(right_image)

        # only keep required keys and prepare for network
        inputs = {'image': inputs['left_image'],
                  'stereo_image': right_image,
                  'disparity': projection_disparity.astype(float),
                  'mono_disparity': inputs['loaded_disparity'].astype(float),
                  }

        # finally crop to feed width
        for key in ['image', 'stereo_image']:
            inputs[key] = inputs[key].crop((0, 0, self.feed_width, self.feed_height))
        for key in ['disparity', 'mono_disparity']:
            inputs[key] = inputs[key][:, :self.feed_width]

        # self.preprocess(inputs)
        return inputs
if __name__=="__main__":
    
    dataset=BaseDataset(r"/home/notebook/data/group/Ruby/stereo-from-mono-master/")
    # 输入数据path
    # save_path=r"F:\code\stereo-from-mono-master\langan_refine_data/"
    save_path=r"/home/notebook/data/group/Ruby/stereo-from-mono-master/refine_results/"
    # 输出数据path
    count=5000 # 原值是5000
    for data in dataset:
        print(count)
        left=np.array(data["image"])
        right=np.array(data["stereo_image"])
        disparity=np.array(data["disparity"])
        mono_disparity=np.array(data["mono_disparity"])
        left=cv2.cvtColor(left,cv2.COLOR_RGB2BGR)
        right=cv2.cvtColor(right,cv2.COLOR_RGB2BGR)
        print(save_path+str(count)+"_left.jpg")
        cv2.imwrite(save_path+str(count)+"_left.jpg",left)
        cv2.imwrite(save_path+str(count)+"_right.jpg",right)
        cv2.imwrite(save_path+str(count)+"_disp.png",disparity*10)
        count+=1

        # plt.subplot(1,4,1)
        # plt.imshow(left)
        # plt.subplot(1,4,2)
        # plt.imshow(right)
        # plt.subplot(1,4,3)
        # plt.imshow(disparity,cmap="gray")
        #
        # plt.subplot(1,4,4)
        # plt.imshow(mono_disparity,cmap="gray")
        # plt.show()