from __future__ import absolute_import, division, print_function

import os
import torch
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
from torchvision import transforms

# Set device to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def pil_loader(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)

        with open(dtxt_path, 'r') as fd:
            for ld in fd.readlines():
                ld = ld.strip('\n')
                all_depth.append(ld)

        self.data_path = data_path
        self.depth_path = all_depth
        self.img_path = all_img
        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.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

        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))

        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.img_path)

    def load_images(self, idx, do_flip=False):
        image = self.loader(self.img_path[idx])

        if do_flip:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)

        frame = random.choice(self.img_path)
        background = self.loader(frame)
        return image, background

    def load_disparity(self, idx, do_flip=False):
        depth = self.loader(self.depth_path[idx]).convert("L")
        disparity = np.squeeze(depth)

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

    def preprocess(self, inputs):
        if self.is_train and random.random() > 0.5:
            # Create a ColorJitter object with the desired parameters
            color_aug = transforms.ColorJitter(
                brightness=self.brightness, 
                contrast=self.contrast, 
                saturation=self.saturation, 
                hue=self.hue
        )

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

        for key in ['image', 'stereo_image']:
            if self.disable_normalisation:
                inputs[key] = self.to_tensor(inputs[key]).to(device)
            else:
                inputs[key] = ((self.to_tensor(inputs[key]) - 0.45) / 0.225).to(device)

        if self.has_gt:
            inputs['disparity'] = torch.from_numpy(inputs['disparity']).float().to(device)

        if inputs.get('mono_disparity') is not None:
            inputs['mono_disparity'] = torch.from_numpy(inputs['mono_disparity']).float().to(device)
    def process_disparity(self, disparity, max_disparity_range=(40, 196)):
        disparity = disparity.copy()

        min_disp = disparity.min()
        if min_disp < 0:
            disparity += np.abs(min_disp)

        disparity /= disparity.max()

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

        if not self.disable_sharpening:
            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):
                pass

        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:
                target_height = height
                target_width = width
            else:
                current_ratio = height / width
                target_ratio = self.feed_height / self.process_width

                if current_ratio < target_ratio:
                    target_height = self.feed_height
                    target_width = int(self.feed_height / height * width)
                elif current_ratio > target_ratio:
                    target_height = int(self.process_width / width * height)
                    target_width = self.process_width
                else:
                    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)

        if target_height == self.feed_height and target_width == self.process_width:
            pass
        else:
            self.crop_all(inputs)

        return inputs

    def crop_all(self, inputs):

        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):

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

        disp = cv2.resize(disp.astype(float), (width, height))
        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]
            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)

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

        mask = self.get_occlusion_mask(pix_locations)
        masked_pix_locations = pix_locations * mask - self.process_width * (1 - mask)

        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
        weights = np.expand_dims(weights, -1)
        warped_image = warped_image[0] * weights[1] + warped_image[1] * weights[0]
        warped_image *= 255.

        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))

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

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

        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

        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

        inputs = self.prepare_sizes(inputs)

        inputs['background'] = transfer_color(np.array(inputs['background']),
                                              np.array(inputs['left_image']))

        inputs['disparity'] = \
            self.process_disparity(inputs['loaded_disparity'],
                                   max_disparity_range=(10, self.max_disparity))

        projection_disparity = inputs['disparity']
        right_image = self.project_image(inputs['left_image'],
                                         projection_disparity, inputs['background'])

        right_image = self.augment_synthetic_image(right_image)

        inputs = {'image': inputs['left_image'],
                  'stereo_image': right_image,
                  'disparity': projection_disparity.astype(float),
                  'mono_disparity': inputs['loaded_disparity'].astype(float),
                  }

        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/")
    save_path=r"/home/notebook/data/group/Ruby/stereo-from-mono-master/refine_results/"
    count=5000 # 原值是5000
    for data in dataset:
        print(count)
        left = np.array(data["image"].cpu().numpy(), dtype=np.uint8)  # Move to CPU and convert to numpy
        right = np.array(data["stereo_image"].cpu().numpy(), dtype=np.uint8)
        disparity = np.array(data["disparity"].cpu().numpy())
        mono_disparity = np.array(data["mono_disparity"].cpu().numpy())
        left = cv2.cvtColor(left, cv2.COLOR_RGB2BGR)
        right = cv2.cvtColor(right, cv2.COLOR_RGB2BGR)
        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