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


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
        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, 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):
    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 + 5000
    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 = '/home/notebook/data/group/Ruby/stereo-from-mono/langan_img.txt'
    dtxt_path = '/home/notebook/data/group/Ruby/stereo-from-mono/langan_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/langan"
    # count_start, count_end = 0, 356
    count_start, count_end = 0, len(all_img)
    
    feed_height, feed_width, process_width = 1024, 768, 768 + 20
    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 range(count_start, count_end)]
        processed_data = [p.get() for p in results]

    print("Processing complete.")