from abc import ABCMeta, abstractmethod

import numpy as np
import tensorflow as tf
from skimage.draw import circle
import Constants as Constants
import datasets.Util.Util as Util
from Log import log
from datasets.Augmentors import parse_augmentors
from datasets.Util.Batch import create_batch_dict
from datasets.Util.Normalization import unnormalize
from datasets.Util.Reader import read_images_from_disk, load_label_default, load_img_default
from datasets.Util.Resize import parse_resize_mode, ResizeMode


class Dataset(object, metaclass=ABCMeta):
    def __init__(self, subset):
        self.summaries = []
        self.subset = subset

    def _get_resize_params(self, subset, input_size_default, resize_mode_default=ResizeMode.Unchanged):
        if subset == "valid":
            subset = "val"
        resize_mode_string = self.config.str("resize_mode_" + subset, "")
        # use same as for train if nothing specified
        if resize_mode_string == "":
            resize_mode_string = self.config.str("resize_mode_train", "")
        if resize_mode_string == "":
            resize_mode = resize_mode_default
        else:
            resize_mode = parse_resize_mode(resize_mode_string)
        input_size = self.config.int_list("input_size_" + subset, [])
        # use same as for train if nothing specified
        if len(input_size) == 0:
            if resize_mode == ResizeMode.Unchanged:
                input_size = [None, None]
            else:
                input_size = self.config.int_list("input_size_train", [])
        if len(input_size) == 0:
            input_size = input_size_default
        if resize_mode == ResizeMode.RandomResize:
            input_size = [None, None]
        return resize_mode, input_size

    def _parse_augmentors_and_shuffle(self):
        if self.subset == "train":
            shuffle = True
            augmentor_strs = self.config.unicode_list("augmentors_train", [])
            augmentors = parse_augmentors(augmentor_strs, self.void_label())
            if len(augmentors) == 0:
                print("warning, no data augmentors used on train", file=log.v1)
        else:
            shuffle = False
            augmentor_strs = self.config.unicode_list("augmentors_val", [])
            augmentors = parse_augmentors(augmentor_strs, self.void_label())
        return augmentors, shuffle

    @abstractmethod
    def num_classes(self):
        pass

    @abstractmethod
    def num_examples_per_epoch(self):
        pass

    # should contain inputs, labels, tags, maybe raw_labels, index_img
    @abstractmethod
    def create_input_tensors_dict(self, batch_size):
        pass

    @abstractmethod
    def void_label(self):
        pass


class ImageDataset(Dataset):
    def __init__(self, dataset_name, default_path, num_classes, config, subset, coord, image_size, void_label=255,
                 fraction=1.0, label_postproc_fn=lambda x: x, label_load_fn=load_label_default,
                 img_load_fn=load_img_default, ignore_classes=[]):
        super(ImageDataset, self).__init__(subset)
        self._num_classes = num_classes
        self._void_label = void_label
        assert subset in ("train", "valid"), subset
        if subset == "train":
            self.epoch_length = config.int("epoch_length", - 1)
        else:
            self.epoch_length = -1
        self.config = config
        self.data_dir = config.str(dataset_name + "_data_dir", default_path)
        self.coord = coord
        self.image_size = image_size
        self.inputfile_lists = None
        self.fraction = fraction
        self.label_postproc_fn = label_postproc_fn
        self.label_load_fn = label_load_fn
        self.img_load_fn = img_load_fn
        self.use_summaries = self.config.bool("use_summaries", False)
        self.ignore_classes = ignore_classes

    def _load_inputfile_lists(self):
        if self.inputfile_lists is not None:
            return
        self.inputfile_lists = self.read_inputfile_lists()
        # make sure all lists have the same length
        assert all([len(l) == len(self.inputfile_lists[0])
                    for l in self.inputfile_lists])
        if self.fraction < 1.0:
            n = int(self.fraction * len(self.inputfile_lists[0]))
            self.inputfile_lists = [l[:n] for l in self.inputfile_lists]

    def create_input_tensors_dict(self, batch_size):
        self._load_inputfile_lists()
        resize_mode, input_size = self._get_resize_params(
            self.subset, self.image_size, ResizeMode.Unchanged)
        augmentors, shuffle = self._parse_augmentors_and_shuffle()

        #inputfile_tensors = [tf.convert_to_tensor(l, dtype=tf.string) for l in self.inputfile_lists]
        inputfile_tensors = [tf.convert_to_tensor(
            l) for l in self.inputfile_lists]
        queue = tf.train.slice_input_producer(
            inputfile_tensors, shuffle=shuffle)

        tensors_dict, summaries = self._read_inputfiles(
            queue, resize_mode, input_size, augmentors)
        tensors_dict, summ = create_batch_dict(batch_size, tensors_dict)
        if summ is not None:
            summaries.append(summ)
        if Constants.BBOXES in tensors_dict and Constants.IDS in tensors_dict and Constants.CLASSES in tensors_dict:
            assert "labels" not in tensors_dict
            tensors_dict["labels"] = (tensors_dict[Constants.BBOXES], tensors_dict[Constants.IDS],
                                      tensors_dict[Constants.CLASSES])
            if Constants.IGNORE_REGIONS in tensors_dict:
                tensors_dict["labels"] += (
                    tensors_dict[Constants.IGNORE_REGIONS],)
                if Constants.SCENE_INFOS in tensors_dict:
                    tensors_dict["labels"] += (
                        tensors_dict[Constants.SCENE_INFOS],)

        if self.use_summaries:
            summaries += self._create_summaries(tensors_dict)

        self.summaries += summaries
        return tensors_dict

    def _create_summaries(self, tensors_dict):
        inputs = tensors_dict["inputs"]
        input_imgs = unnormalize(inputs[:, :, :, :3])
        # count is incremented after each summary creation. This helps us to keep track of the number of channels
        start = 3
        summaries = []
        # Old label would be present either as a single channel, or along with 2 other distance transform channels.
        if inputs.get_shape()[-1] in [4, 6, 8]:
            if Constants.OLD_LABEL_AS_DT in tensors_dict:
                [input_imgs] = tf.py_func(self.add_clicks, [tf.concat([input_imgs, inputs[:, :, :, start:start + 1]],
                                                                      axis=3), 'g'], [tf.float32])
            else:
                [old_label, input_imgs] = tf.py_func(self.get_masked_image, [input_imgs, inputs[:, :, :, start:start + 1]],
                                                     [tf.float32, tf.float32])

            summ = tf.summary.image(
                "old_labels", inputs[:, :, :, start:start + 1])
            summaries.append(summ)
            start += 1

        # Add clicks to the image so that they can be viewed in tensorboard
        if inputs.get_shape()[-1] > 6:
            [input_imgs] = tf.py_func(self.add_clicks, [tf.concat([input_imgs, inputs[:, :, :, start:start + 1]],
                                                                  axis=3), 'r'], [tf.float32])
            [input_imgs] = tf.py_func(self.add_clicks, [tf.concat([input_imgs, inputs[:, :, :, start + 1:start + 2]],
                                                                  axis=3), 'r'], [tf.float32])
            [input_imgs] = tf.py_func(self.add_clicks, [tf.concat([input_imgs, inputs[:, :, :, start + 2:start + 3]],
                                                                  axis=3), 'g'], [tf.float32])
            [input_imgs] = tf.py_func(self.add_clicks, [tf.concat([input_imgs, inputs[:, :, :, start + 3:start + 4]],
                                                                  axis=3), 'g'], [tf.float32])
        elif inputs.get_shape()[-1] > 4:
            [input_imgs] = tf.py_func(self.add_clicks, [tf.concat([input_imgs, inputs[:, :, :, start:start + 1]],
                                                                  axis=3), 'r'], [tf.float32])
            [input_imgs] = tf.py_func(self.add_clicks, [tf.concat([input_imgs, inputs[:, :, :, start + 1:start + 2]],
                                                                  axis=3), 'g'], [tf.float32])

        # bounding boxes
        if Constants.BBOXES in tensors_dict:
            bboxes = tensors_dict[Constants.BBOXES]
            # permute y1, y2, x1, x2 -> y1, x1, y2, x1
            bboxes = tf.stack([bboxes[..., 0], bboxes[..., 2],
                               bboxes[..., 1], bboxes[..., 3]], axis=-1)
            # normalize bboxes to [0..1]
            height = tf.shape(input_imgs)[1]
            width = tf.shape(input_imgs)[2]
            bboxes = tf.cast(bboxes, tf.float32) / tf.cast(
                tf.stack([height, width, height, width], axis=0), tf.float32)
            imgs_with_bboxes = tf.image.draw_bounding_boxes(input_imgs, bboxes)
            summ = tf.summary.image("inputs", imgs_with_bboxes)
            summaries.append(summ)
        else:
            summ = tf.summary.image("inputs", input_imgs)
            summaries.append(summ)

        if not isinstance(tensors_dict["labels"], tuple):
            # will only work well for binary segmentation
            summ = tf.summary.image("labels", tensors_dict["labels"] * 255)
            summaries.append(summ)

        # Append the distance transforms, if they are available.
        if inputs.get_shape()[-1] > 6:
            # Get negative distance transform from the extra input channels.
            summ = tf.summary.image(
                Constants.DT_NEG, inputs[:, :, :, start:start + 1])
            summaries.append(summ)
            summ = tf.summary.image(
                Constants.DT_NEG, inputs[:, :, :, start+1:start + 2])
            summaries.append(summ)
            summ = tf.summary.image(
                Constants.DT_POS, inputs[:, :, :, start + 2:start + 3])
            summaries.append(summ)
            summ = tf.summary.image(
                Constants.DT_POS, inputs[:, :, :, start + 3:start + 4])
            summaries.append(summ)
        elif inputs.get_shape()[-1] > 4:
            # Get negative distance transform from the extra input channels.
            summ = tf.summary.image(
                Constants.DT_NEG, inputs[:, :, :, start:start + 1])
            summaries.append(summ)
            summ = tf.summary.image(
                Constants.DT_POS, inputs[:, :, :, start + 1:start + 2])
            summaries.append(summ)

        return summaries

    # default implementation, should in many cases be overwritten
    def _read_inputfiles(self, queue, resize_mode, input_size, augmentors):
        tensors, summaries = read_images_from_disk(queue, input_size, resize_mode, label_postproc_fn=self.label_postproc_fn,
                                                   augmentors=augmentors, label_load_fn=self.label_load_fn,
                                                   img_load_fn=self.img_load_fn)
        return tensors, summaries

    def add_clicks(self, inputs, c):
        out_imgs = None

        for input in inputs:
            img = input[:, :, :3]
            # Radius of the point to be diplayed
            r = 3
            pts = np.where(input[:, :, 3] == 0)
            pts_zipped = list(zip(pts[0], pts[1]))
            if len(pts[0]) > 0:
                for pt in pts_zipped:
                    if r < pt[0] < img.shape[0] - r and r < pt[1] < img.shape[1] - r:
                        rr, cc = circle(pt[0], pt[1], 5, img.shape)
                        img[rr, cc, :] = [np.max(img), np.min(img), np.min(img)] if c == 'r' \
                            else [np.min(img), np.min(img), np.max(img)]

            img = img[np.newaxis, :, :, :]
            if out_imgs is None:
                out_imgs = img
            else:
                out_imgs = np.concatenate((out_imgs, img), axis=0)

        return out_imgs.astype(np.float32)

    def get_masked_image(self, inputs, old_labels):
        out_imgs = None
        out_labels = None
        for input_img, old_label_unmodified in zip(inputs, old_labels):
            if len(np.unique(old_label_unmodified)) > 2:
                old_label = np.where(old_label_unmodified*255 <= 128, 1, 0)
            else:
                old_label = old_label_unmodified

            out_img = Util.get_masked_image(input_img, old_label[:, :, 0])
            out_img = out_img[np.newaxis, :, :, :]
            out_label = old_label[np.newaxis, :, :, :]
            if out_imgs is None:
                out_imgs = out_img
            else:
                out_imgs = np.concatenate((out_imgs, out_img), axis=0)

            if out_labels is None:
                out_labels = out_label
            else:
                out_labels = np.concatenate((out_labels, out_label), axis=0)

        return out_labels.astype(np.float32), out_imgs.astype(np.float32)

    @abstractmethod
    def read_inputfile_lists(self):
        pass

    def num_examples_per_epoch(self):
        if self.epoch_length != -1:
            return self.epoch_length
        else:
            self._load_inputfile_lists()
            return len(self.inputfile_lists[0])

    def num_classes(self):
        return self._num_classes

    def void_label(self):
        return self._void_label

    def ignore_classes(self):
        return self.ignore_classes
