import os

import numpy as np
import tensorflow as tf
from random import randint
from math import sqrt, ceil

import cv2

def image_scaling(img, label, input_size):
    """
    Randomly scales the images between 0.4 to 1.0 times the original size.

    Args:
      img: Training image to scale.
      label: Segmentation mask to scale.
    """

    scale = tf.random_uniform([1], minval = 0.5, maxval = 1.3, dtype=tf.float32, seed=None)
    h_new = tf.to_int32(tf.multiply(tf.to_float(tf.shape(img)[0]), scale))
    w_new = tf.to_int32(tf.multiply(tf.to_float(tf.shape(img)[1]), scale))
    new_shape = tf.squeeze(tf.stack([h_new, w_new]), squeeze_dims=[1])

    img = tf.image.resize_images(img, new_shape)
    label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), new_shape)
    label = tf.squeeze(label, squeeze_dims=[0])

    img = tf.image.resize_image_with_crop_or_pad(img, input_size[0], input_size[1])
    label = tf.image.resize_image_with_crop_or_pad(label, input_size[0], input_size[1])


    return img, label

def image_mirroring(img, label):
    """
    Randomly mirrors the images.

    Args:
      img: Training image to mirror.
      label: Segmentation mask to mirror.
    """

    distort_left_right_random = tf.random_uniform([1], 0, 1.0, dtype=tf.float32)[0]
    mirror = tf.less(tf.stack([1.0, distort_left_right_random, 1.0]), 0.5)
    mirror = tf.boolean_mask([0, 1, 2], mirror)
    img = tf.reverse(img, mirror)
    label = tf.reverse(label, mirror)
    return img, label

def random_roi(img, answer, y_min, y_max, x_left_gt, x_right_gt, shape, roi_params):
    """
    Randomly mirrors the images.

    Args:
      img: Training image.
      answer: Segmentation mask.
      y_min, y_max - y roi values
      x_left_gt, x_right_gt - x bounds of grountruth object's location, from this values will be random indenting
      shape: Shape of image
      roi_params: max and min available indenting
    """
    max_x_indent, min_x_indent = roi_params

    init_w = shape[0]

    min_indent = tf.constant(min_x_indent,dtype=tf.int32)
    max_indent = tf.constant(max_x_indent,dtype=tf.int32)

    tf_zero = tf.constant(0, dtype=tf.int32)

    x_left  = tf.random_uniform(shape = [1], minval = tf.maximum(tf.subtract(x_left_gt, max_indent), tf_zero), maxval = tf.maximum(tf.subtract(x_left_gt, min_indent), tf_zero), dtype = tf.int32)[0]
    x_right = tf.random_uniform(shape = [1], minval = tf.minimum(tf.add(x_right_gt, min_indent), init_w), maxval = tf.minimum(tf.add(x_right_gt, max_indent), init_w), dtype = tf.int32)[0]

    y_top = y_min
    y_bot = y_max

    image = tf.slice(img, [y_top, x_left, 0], [tf.subtract(y_bot,y_top), tf.subtract(x_right,x_left), 3])
    label = tf.slice(answer, [y_top, x_left, 0], [tf.subtract(y_bot,y_top), tf.subtract(x_right,x_left), 1])

    new_shape = tf.convert_to_tensor([tf.shape(image)[1], tf.shape(image)[0]], dtype = tf.int32)

    return image, label, new_shape

def get_rois(max_x_indent, min_x_indent, num_rois):
    '''
    Get num_rois number of x indents
    '''
    assert(num_rois>1),"Number of rois for validation mist be greater then 1!"
    if (num_rois == 2):
        return [(min_x_indent, min_x_indent), (max_x_indent, max_x_indent)]

    num_of_values = ceil(sqrt(num_rois))-1

    step = (max_x_indent - min_x_indent)/num_of_values

    step = ceil(step)
    val_indents = [i for i in range(min_x_indent, max_x_indent, step)]
    val_indents.append(max_x_indent)
    val_roi_indents = []

    for indent in val_indents:
        for indent_ in val_indents:
            val_roi_indents.append((indent, indent_))

            if len(val_roi_indents) == num_rois:
                return val_roi_indents

    return val_roi_indents

def fixed_roi(img, answer, y_min, y_max, x_left_gt, x_right_gt, shape, roi_params, num_of_valid_step, num_rois):
    '''
    Apply on of fixed roi(depends on num_of_valid_step value) during validation
    '''
    init_w = shape[0]

    max_x_indent,min_x_indent = roi_params
    val_roi_indents = get_rois(max_x_indent, min_x_indent, num_rois)
    #print("\nROI indents==========>",val_roi_indents[num_of_valid_step])

    x_left_indent = tf.constant(val_roi_indents[num_of_valid_step][0],dtype=tf.int32)
    x_right_indent = tf.constant(val_roi_indents[num_of_valid_step][1],dtype=tf.int32)

    tf_zero = tf.constant(0, dtype=tf.int32)

    x_left  = tf.math.maximum(tf.subtract(x_left_gt, x_left_indent), tf_zero)
    x_right = tf.math.minimum(tf.math.add(x_right_gt, x_right_indent), init_w)

    y_top = y_min
    y_bot = y_max

    image = tf.slice(img, [y_top, x_left, 0], [tf.subtract(y_bot,y_top), tf.subtract(x_right,x_left), 3])
    label = tf.slice(answer, [y_top, x_left, 0], [tf.subtract(y_bot,y_top), tf.subtract(x_right,x_left), 1])

    new_shape = tf.convert_to_tensor([tf.shape(image)[1], tf.shape(image)[0]], dtype = tf.int32)

    return image, label, new_shape

def random_crop_and_pad_image_and_labels(image, label, crop_h, crop_w, pad_h, pad_w, input_size, CROP_PROB, PAD_PROB, ignore_label):
    """
    Randomly crop and pads the input images.

    Args:
      image: Training image to crop/ pad.
      label: Segmentation mask to crop/ pad.
      crop_h: Height of cropped segment.
      crop_w: Width of cropped segment.
    """

    label = tf.cast(label, dtype = tf.float32)
    label = label - ignore_label # Needs to be subtracted and later added due to 0 padding.
    combined = tf.concat(axis = 2, values = [image, label])

    last_image_dim = tf.shape(image)[-1]
    #last_label_dim = tf.shape(label)[-1]

    use_crop = tf.random_uniform(shape = [1], minval = 0.0, maxval = 1.0, dtype = tf.float32)[0]
    use_pad = tf.random_uniform(shape = [1], minval = 0.0, maxval = 1.0, dtype = tf.float32)[0]
    #use_zoom = tf.random_uniform(shape = [1], minval = 0.0, maxval = 1.0, dtype = tf.float32)[0]
    #focal_len = tf.random_uniform(shape = [1], minval = MIN_FOCAL, maxval = MAX_FOCAL, dtype = tf.float32)[0]
    crop_size = tf.stack([crop_h, crop_w, 4], axis = 0)

    rows_shift = tf.cast((pad_h - input_size[1]) / 2, dtype = tf.int32)
    cols_shift = tf.cast((pad_w - input_size[0]) / 2, dtype = tf.int32)
    pads = tf.stack([[rows_shift, rows_shift], [cols_shift, cols_shift], [0, 0]])

    combined_crop = tf.cond(use_crop > CROP_PROB, lambda : combined, lambda : tf.random_crop(combined, crop_size))
    combined_crop = tf.cond(use_pad > PAD_PROB, lambda : combined_crop, lambda : tf.pad(combined_crop, pads))

    img_crop = combined_crop[:, :, :last_image_dim]
    label_crop = combined_crop[:, :, last_image_dim:]
    label_crop = label_crop + ignore_label
    label_crop = tf.cast(label_crop, dtype = tf.uint8)

    return img_crop, label_crop

def random_rotation(image, label, rot_params):
    ANGLE_RADIANS, ROT_PROB = rot_params
    # max rotation angle ~ 13 degrees
    use_rotation = tf.random_uniform(shape = [1], minval = 0.0, maxval = 1.0, dtype = tf.float32)[0]
    rotation_radians = tf.random_uniform([], minval = -ANGLE_RADIANS, maxval = ANGLE_RADIANS, dtype=tf.float32, seed=None)

    img = tf.cond(use_rotation > ROT_PROB, lambda : image, lambda : tf.contrib.image.rotate(image, rotation_radians, interpolation='BILINEAR'))
    lbl = tf.cond(use_rotation > ROT_PROB, lambda : label, lambda : tf.contrib.image.rotate(label, rotation_radians, interpolation='BILINEAR'))

    return img, lbl

def read_labeled_image_list(data_list, input_shape):
    """Reads txt file containing paths to images and ground truth masks.

    Args:
      data_list: path to the file with lines of the form '/path/to/image /path/to/mask real_dx real_dy min_y max_y left_x right_x'.

    Returns:
      Lists with:
            all file names for images
            all file names for masks
            all x-pixel's sizes for each image
            all y pixel's sizes for each image
            all initial shapes of images
            all top y coordinates of roi for each image
            all bottom y coordinates of roi for each image
            all max left x coordinates of roi for each image
            all min right x coordinates of roi for each image
    """
    f = open(data_list, 'r')
    images = []
    masks = []

    # x-pixel's size in cm
    PhysicalDeltaXs = []
    # y pixel's size in cm
    PhysicalDeltaYs= []

    # ys for roi
    yMins = []
    yMaxs = []

    # x bounds of gt bbox (around gt points)
    leftXs = []
    rightXs = []

    init_shapes = []

    for line in f:
        real_dx = None
        real_dy = None

        # Roi top coordinate
        min_y = None
        # Roi bottom coordinate
        max_y = None

        # Maximum left x coordinate of roi (need for random roi selection)
        left_x = None
        # minimum right x coordinate of roi (need for random roi selection)
        right_x = None

        try:
            data = line.strip("\n").split(' ')
            image, mask, real_dx, real_dy, min_y, max_y, left_x, right_x = data

        except ValueError:
            image = mask = line.strip("\n")

        images.append(image)
        masks.append(mask)

        yMins.append(int(min_y))
        yMaxs.append(int(max_y))

        leftXs.append(int(left_x))
        rightXs.append(int(right_x))

        if not real_dx is None:
            PhysicalDeltaX = float(real_dx)#pixel size in cm
            PhysicalDeltaX = PhysicalDeltaX if PhysicalDeltaX > 0 else -1 * PhysicalDeltaX
            PhysicalDeltaXs.append(PhysicalDeltaX)

        if not real_dy is None:
            PhysicalDeltaY = float(real_dy)#pixel size in cm
            PhysicalDeltaY = PhysicalDeltaY if PhysicalDeltaY > 0 else -1 * PhysicalDeltaY
            PhysicalDeltaYs.append(PhysicalDeltaY)

        im = cv2.imread(image)
        #if im is None:
        #    print(image)
        init_h, init_w, _ = im.shape
        init_shapes.append([init_w, init_h])

    return images, masks, PhysicalDeltaXs, PhysicalDeltaYs, init_shapes, yMins, yMaxs, leftXs, rightXs

def read_images_from_disk(input_queue, input_size,
                random_scale, random_mirror, random_rotate, random_roi_,
                img_mean, roi_params, crop_params, pad_params, rot_params, ignore_label,
                train = True, num_of_valid_step = 0, num_rois = 12, gray = False):
    """Read one image and its corresponding mask with optional pre-processing.

    Args:
      input_queue: tf queue with paths to the image and its mask.
      input_size: a tuple with (height, width) values.
                  If not given, return images of original size.
      random_scale: whether to randomly scale the images prior
                    to random crop.
      random_mirror: whether to randomly mirror the images prior
                    to random crop.
      random_rotate: whether to randomly rotate the images prior
                    to random rotation.
      random_roi_: whether to apply random rois to images

      img_mean: vector of mean colour values.
      roi_params: maximum available and minimum available x-axis indents during calculating random roi
    Returns:

    """
    img_contents = tf.read_file(input_queue[0])
    label_contents = tf.read_file(input_queue[1])

    img = tf.image.decode_jpeg(img_contents, channels=3)
    answer = tf.image.decode_png(label_contents, channels=1)

    ## select random roi and crop
    new_init_shape = input_queue[4]
    if random_roi_:

        init_shape = input_queue[4]

        y_min = input_queue[5]
        y_max = input_queue[6]

        x_left = input_queue[7]
        x_right = input_queue[8]

        if train:
            img, answer, new_init_shape = random_roi(img, answer, y_min, y_max, x_left, x_right, init_shape, roi_params)
        else:
            img, answer, new_init_shape = fixed_roi(img, answer, y_min, y_max, x_left, x_right, init_shape, roi_params, num_of_valid_step, num_rois)

    if train and random_rotate:
        img, answer = random_rotation(img, answer, rot_params)

    img = tf.image.resize_images(img, input_size)
    answer = tf.image.resize_nearest_neighbor(tf.expand_dims(answer, 0), input_size)
    answer = tf.squeeze(answer, squeeze_dims=[0])

    # converting to bgr
    #img_r, img_g, img_b = tf.split(axis=2, num_or_size_splits=3, value=img)
    #img = tf.cast(tf.concat(axis=2, values=[img_b, img_g, img_r]), dtype=tf.float32)

    # Extract mean.
    if not gray:
        image = img - img_mean
    else:
        image = img

    if input_size is None:
        print('WTF?!?!')
        quit()

    h, w = input_size

    # Randomly scale the images and labels.
    if train and random_scale:
        image, answer = image_scaling(image, answer, input_size)

    # Randomly mirror the images and labels.
    if train and random_mirror:
        image, answer = image_mirroring(image, answer)

    if train:
        MIN_CROP, MAX_CROP, CROP_PROB = crop_params
        h_rate = tf.random_uniform(shape = [1], minval = MIN_CROP, maxval = MAX_CROP, dtype = tf.float32)
        w_rate = tf.random_uniform(shape = [1], minval = MIN_CROP, maxval = MAX_CROP, dtype = tf.float32)
        crop_h = tf.cast(tf.cast(h, tf.float32) * h_rate, tf.int32)[0]
        crop_w = tf.cast(tf.cast(w, tf.float32) * w_rate, tf.int32)[0]

        MIN_PAD, MAX_PAD, PAD_PROB = pad_params
        h_rate = tf.random_uniform(shape = [1], minval = MIN_PAD, maxval = MAX_PAD, dtype = tf.float32)
        w_rate = tf.random_uniform(shape = [1], minval = MIN_PAD, maxval = MAX_PAD, dtype = tf.float32)
        pad_h = tf.cast(tf.cast(h, tf.float32) * h_rate, tf.int32)[0]
        pad_w = tf.cast(tf.cast(w, tf.float32) * w_rate, tf.int32)[0]

        img, label = random_crop_and_pad_image_and_labels(image, answer, crop_h, crop_w, pad_h, pad_w, input_size, CROP_PROB, PAD_PROB, ignore_label)

    else:
        img = image
        label = answer

    img = tf.image.resize_images(img, input_size)
    if gray:
        img = tf.image.rgb_to_grayscale(img)

    label = tf.image.resize_nearest_neighbor(tf.expand_dims(label, 0), input_size)
    label = tf.squeeze(label, squeeze_dims=[0])
    if gray:
        img.set_shape([h, w, 1])
    else:
        img.set_shape([h, w, 3])
    label.set_shape([h, w, 1])

    return img, label, input_queue[2], input_queue[3], new_init_shape

class ImageReader(object):
    '''Generic ImageReader which reads images and corresponding segmentation
       masks from the disk, and enqueues them into a TensorFlow queue.
    '''
    def __init__(self, coord, params, train = True, num_of_valid_step = 0):
        max_x_indent = params["MAX_X_INDENT"]
        min_x_indent = params["MIN_X_INDENT"]
        self.is_gray = params["IS_GRAY"]
        self.data_list = params["DATA_LIST"]
        self.input_size = params["INPUT_SIZE"]
        self.img_mean = params["IMG_MEAN"]
        self.ignore_label = params["IGNORE_LABEL"]
        self.random_roi = params["RANDOM_ROI"]

        if train:
            self.random_rotate = params["RANDOM_ROTATE"]
            self.random_scale = params["RANDOM_SCALE"]
            self.random_mirror = params["RANDOM_MIRROR"]
            self.num_rois = 0
        else:
            self.random_rotate, self.random_scale, self.random_mirror = [False]*3
            self.num_rois = params["NUM_ROIS"]

        self.coord = coord
        self.train = train
        self.num_of_valid_step = num_of_valid_step

        self.roi_params = [max_x_indent, min_x_indent]
        if train:
            self.crop_params = [params["MIN_CROP"], params["MAX_CROP"], params["CROP_PROB"]]
            self.pad_params = [params["MIN_PAD"], params["MAX_PAD"], params["PAD_PROB"]]
            self.rot_params = [params["ANGLE_RADIANS"], params["ROT_PROB"]]
        else:
            self.crop_params = []
            self.pad_params = []
            self.rot_params = []


        self.image_list, self.label_list, self.x_pixels_data, self.y_pixels_data, self.init_shapes, \
                        self.yMins, self.yMaxs, self.leftXs, self.rightXs = read_labeled_image_list(self.data_list, self.input_size)

        self.images = tf.convert_to_tensor(self.image_list, dtype=tf.string)
        self.labels = tf.convert_to_tensor(self.label_list, dtype=tf.string)

        self.y_min_data = tf.convert_to_tensor(self.yMins, dtype=tf.int32)
        self.y_max_data = tf.convert_to_tensor(self.yMaxs, dtype=tf.int32)

        self.x_left_data = tf.convert_to_tensor(self.leftXs, dtype=tf.int32)
        self.x_right_data = tf.convert_to_tensor(self.rightXs, dtype=tf.int32)

        self.x_datas = tf.convert_to_tensor(self.x_pixels_data, dtype=tf.float32)
        self.y_datas = tf.convert_to_tensor(self.y_pixels_data, dtype=tf.float32)

        self.shapes = tf.convert_to_tensor(self.init_shapes, dtype=tf.int32)

        self.queue = tf.train.slice_input_producer([self.images, self.labels, self.x_datas, self.y_datas, self.shapes,
                                                    self.y_min_data, self.y_max_data, self.x_left_data, self.x_right_data],
                                                   shuffle=self.train) # not shuffling if it is val
        self.image, self.label, self.x_data, self.y_data, self.shape = read_images_from_disk(self.queue, self.input_size,
                                                                                                self.random_scale, self.random_mirror, self.random_rotate, self.random_roi,
                                                                                                self.img_mean, self.roi_params, self.crop_params, self.pad_params, self.rot_params,
                                                                                                self.ignore_label, self.train, self.num_of_valid_step, self.num_rois, self.is_gray)

    def dequeue(self, num_elements):
        '''Pack images, labels, real x-pixel's sizes, real y-pixel's sizes and images shapes before applying roi (need for calculating in cm)

        Args:
          num_elements: the batch size.'''
        image_batch, label_batch, x_data_batch, y_data_batch, shapes_batch = tf.train.batch([self.image, self.label, self.x_data, self.y_data, self.shape],
                                                                                            num_elements)

        return image_batch, label_batch, x_data_batch, y_data_batch, shapes_batch