from __future__ import absolute_import, division, print_function

import os
import random
import numpy as np
from PIL import Image
from skimage.filters import gaussian, sobel
from scipy.interpolate import griddata
import cv2
import torch
from torchvision import transforms
import multiprocessing as mp
from utils import transfer_color
import matplotlib.pyplot as plt
import sys
from tqdm import tqdm


def pil_loader(path):
    # print(path)
    with open(path, 'rb') as f:
        from PIL import Image
        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)))


def load_images(img_path, idx, do_flip=False):
    image = pil_loader(img_path[idx])
    if do_flip:
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
    frame = random.choice(img_path)
    background = pil_loader(frame)
    return image, background


def load_disparity(depth_path, idx, do_flip=False):
    depth = pil_loader(depth_path[idx]).convert("L")
    disparity = np.squeeze(np.array(depth))
    if do_flip:
        disparity = disparity[:, ::-1]
    return disparity


def preprocess(inputs, feed_width, feed_height, disable_normalisation, has_gt, to_tensor, is_train, brightness, contrast, saturation, hue):
    if is_train and random.random() > 0.5:
        color_aug = transforms.ColorJitter.get_params(
            brightness, contrast, saturation, hue)
        for key in ['image', 'stereo_image']:
            inputs[key] = color_aug(inputs[key])

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

    if 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(disparity, feed_height, process_width, max_disparity_range=(40, 196), disable_sharpening=False):
    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 disable_sharpening:
        edges = sobel(disparity) > 3
        disparity[edges] = 0
        mask = disparity > 0
        ys, xs = np.meshgrid(np.arange(process_width), np.arange(feed_height))
        try:
            disparity = griddata(np.stack([ys[mask].ravel(), xs[mask].ravel()], 1),
                                 disparity[mask].ravel(), np.stack([ys.ravel(),
                                                                    xs.ravel()], 1),
                                 method='nearest').reshape(feed_height, process_width)
        except (ValueError, IndexError):
            pass
    return disparity


def prepare_sizes(inputs, feed_height, process_width, keep_aspect_ratio):
    height, width, _ = np.array(inputs['left_image']).shape
    # if keep_aspect_ratio:
    #     # if feed_height <= height and process_width <= width:
    #     #     target_height, target_width = height, width
    #     if feed_height <= height and process_width <= width:
    #         target_height, target_width = feed_height, feed_width
    #         ### 这边改一下
    #     else:
    #         current_ratio = height / width
    #         target_ratio = feed_height / process_width
    #         if current_ratio < target_ratio:
    #             target_height = feed_height
    #             target_width = int(feed_height / height * width)
    #         elif current_ratio > target_ratio:
    #             target_height = int(process_width / width * height)
    #             target_width = process_width
    #         else:
    #             target_height, target_width = feed_height, process_width
    # else:
    #     target_height, target_width = feed_height, process_width

    
    inputs = resize_all(inputs, feed_height, feed_width)
    
    # inputs = resize_all(inputs, target_height, target_width)
    # if not (target_height == feed_height and target_width == process_width):
    #     inputs = crop_all(inputs, feed_height, process_width)

    return inputs


# def crop_all(inputs, feed_height, process_width):
#     ###要么改这，要么改上面
#     ##改这的话可以多crop几张
#     height, width, _ = np.array(inputs['left_image']).shape
#     top = int(random.random() * (height - feed_height))
#     left = int(random.random() * (width - process_width))
#     right, bottom = left + process_width, top + 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


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(shifted, process_width):
    mask_up, mask_down = shifted > 0, shifted > 0
    shifted_up, shifted_down = np.ceil(shifted), np.floor(shifted)
    for col in range(process_width - 2):
        loc = shifted[:, col:col + 1]
        loc_up, loc_down = np.ceil(loc), 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_down[:, col] = mask_up[:, col] * _mask_up, mask_down[:, col] * _mask_down
    return mask_up + mask_down


def project_image(np_xs, np_ys, feed_height, process_width, image, disp_map, background_image, disable_background):
    image, background_image = np.array(image), np.array(background_image)
    warped_image = np.zeros_like(image).astype(float)
    warped_image = np.stack([warped_image] * 2, 0)
    pix_locations = np_xs - disp_map
    mask = get_occlusion_mask(pix_locations, process_width)
    masked_pix_locations = pix_locations * mask - process_width * (1 - mask)
    weights = np.ones((2, feed_height, process_width)) * 10000

    for col in range(process_width - 1, -1, -1):
        loc = masked_pix_locations[:, col]
        loc_up, loc_down = np.ceil(loc).astype(int), np.floor(loc).astype(int)
        weight_up, weight_down = loc_up - loc, 1 - (loc_up - loc)
        mask = loc_up >= 0
        mask[mask] = weights[0, np.arange(feed_height)[mask], loc_up[mask]] > weight_up[mask]
        weights[0, np.arange(feed_height)[mask], loc_up[mask]] = weight_up[mask]
        warped_image[0, np.arange(feed_height)[mask], loc_up[mask]] = image[:, col][mask] / 255.
        mask = loc_down >= 0
        mask[mask] = weights[1, np.arange(feed_height)[mask], loc_down[mask]] > weight_down[mask]
        weights[1, np.arange(feed_height)[mask], loc_down[mask]] = weight_down[mask]
        warped_image[1, np.arange(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 disable_background:
        occluded = warped_image.max(-1) == 0
        warped_image[occluded] = background_image[occluded]

    return warped_image.astype(np.uint8)


def augment_synthetic_image(disable_synthetic_augmentation, image):
    if disable_synthetic_augmentation:
        return Image.fromarray(image.astype(np.uint8))

    noise = np.random.randn(feed_height, 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)
    return Image.fromarray(image.astype(np.uint8))


def process_and_save(data, save_path, np_xs, np_ys, feed_height, process_width, count):
    left, right_image, disparity, mono_disparity = data
    idx = count
    print('idx=',idx)
    val_path = save_path + str(idx) + "_right.jpg"

    if os.path.isfile(val_path):
        sys.exit()
    else:
        left = cv2.cvtColor(left, cv2.COLOR_RGB2BGR)
        right_image = cv2.cvtColor(right_image, cv2.COLOR_RGB2BGR)
        cv2.imwrite(save_path + str(idx) + "_left.jpg", left)
        cv2.imwrite(save_path + str(idx) + "_right.jpg", right_image)
        cv2.imwrite(save_path + str(idx) + "_disp.png", disparity * 10)

    data = {'left': left, 'right': right_image, 'disparity': disparity, 'mono_disparity': mono_disparity}
    return data


def worker(save_path, data_paths, max_disparity, disable_background, disable_sharpening,
           disable_synthetic_augmentation, process_width, feed_height, feed_width,
           disable_normalisation, has_gt, to_tensor, is_train,
           brightness, contrast, saturation, hue, np_xs, np_ys, keep_aspect_ratio, count):

    img_path, depth_path = data_paths
    do_flip = is_train and random.random() > 0.5

    left_image, background_image = load_images(img_path, count, do_flip)
    loaded_disparity = load_disparity(depth_path, count, do_flip)
    inputs = {'left_image': left_image, 'background': background_image,
              'loaded_disparity': loaded_disparity}

    inputs = prepare_sizes(inputs, feed_height, process_width, keep_aspect_ratio)
    inputs['background'] = transfer_color(np.array(inputs['background']), np.array(inputs['left_image']))

    disparity = process_disparity(inputs['loaded_disparity'], feed_height, process_width, (10, max_disparity))
    right_image = project_image(np_xs, np_ys, feed_height, process_width,
                                inputs['left_image'], disparity, inputs['background'], disable_background)
    right_image = augment_synthetic_image(disable_synthetic_augmentation, right_image)

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

    for key in ['image', 'stereo_image']:
        inputs[key] = inputs[key].crop((0, 0, feed_width, feed_height))

    for key in ['disparity', 'mono_disparity']:
        inputs[key] = inputs[key][:, :feed_width]

    return process_and_save([np.array(inputs["image"]), np.array(inputs["stereo_image"]), 
                             np.array(inputs["disparity"]), np.array(inputs["mono_disparity"])], 
                            save_path, np_xs, np_ys, feed_height, process_width, count)


if __name__ == "__main__":
    # data_path = "/home/notebook/data/group/Ruby/langan_datasets/nightlangan3_deblur/"
    np.random.seed(42)
    
    itxt_path = 'humanwmsk_image.txt'
    dtxt_path = 'humanwmsk_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)

    img_path = all_img
    depth_path = all_depth
    
    save_path = r"/home/notebook/data/group/Ruby/langan_datasets/humanwithmask/stereo4mono/"
    count_start, count_end = 3979, len(all_img)
    
    # feed_height, feed_width, process_width = 1024, 768, 768 + 20
    feed_height, feed_width, process_width = 1024, 768, 768
    disable_background, is_train = False, True
    disable_normalisation, has_gt, disable_sharpening = True, True, True
    max_disparity = 20
    disable_synthetic_augmentation = True
    
    brightness, contrast, saturation, hue = (0.8, 1.2), (0.8, 1.2), (0.8, 1.2), (-0.1, 0.1)
    np_xs, np_ys = np.meshgrid(np.arange(process_width), np.arange(feed_height))
    to_tensor = transforms.ToTensor()
    keep_aspect_ratio = True
    
    data_paths = (img_path, depth_path)
    num_workers = 16
    
    with mp.Pool(processes=num_workers) as pool:
        results = [
            pool.apply_async(worker,
                             args=(save_path, data_paths, max_disparity,
                                   disable_background, disable_sharpening,
                                   disable_synthetic_augmentation, process_width,
                                   feed_height, feed_width, disable_normalisation,
                                   has_gt, to_tensor, is_train, brightness, contrast,
                                   saturation, hue, np_xs, np_ys, keep_aspect_ratio, count))
            for count in tqdm(range(count_start, count_end), desc='Processing Images')
        ]
        processed_data = [p.get() for p in results]

    print("Processing complete.")
   在不改变代码的基本功能和结构的同时，把这个代码改成从count的最大值开始，依次倒着对图像进行处理