import numpy as np

class Dataset(object):
    def __init__(self, params):
        self.width  = params["width"]
        self.height = params["height"]
        self.chns   = params["chns"]

        self.data = {"image" : None, "edges" : None}

        self.data["image"] = np.fromfile(params["images_filepath"], dtype = np.uint8).reshape([-1, self.height, self.width, self.chns])
        self.data["edges"] = np.fromfile(params["edges_filepath"], dtype = np.uint8).reshape([-1, self.height, self.width, 1])
        self.data_source = {"image" : params["images_filepath"], "edges" : params["edges_filepath"]}
        assert self.data["image"].shape[0] == self.data["edges"].shape[0]
        if (params["remove_empty_cell"]):
            not_empty_cells = np.sum(self.data["edges"], axis = (1, 2, 3)) > 0
            self.data["image"] = self.data["image"][not_empty_cells]
            self.data["edges"] = self.data["edges"][not_empty_cells]
        # split into train and validate parts
        validate_cnt = params["validate_batchs"] * params["batch_size"]
        self.split_data(validate_cnt, params["data_split_seed"])
        # data augmentation
        if(params["data_augmentation"]):
            self.augment_data()

    def split_data(self, validate_cnt, random_state):
        train_cnt = self.data["image"].shape[0] - validate_cnt
        p = np.random.RandomState(seed=random_state).permutation(self.data["image"].shape[0])
        self.train_data = {"image" : None, "edges" : None}
        self.train_data["image"] = self.data["image"][p[0:train_cnt]]
        self.train_data["edges"] = self.data["edges"][p[0:train_cnt]]

        self.validate_data = {"image" : None, "edges" : None}
        self.validate_data["image"] = self.data["image"][p[train_cnt:]]
        self.validate_data["edges"] = self.data["edges"][p[train_cnt:]]

        del self.data

    def expand_data(self, data):
        data = np.concatenate([data, data[:,:,::-1,:]], axis = 0)
        data = np.concatenate([data,
                               np.rot90(data, k=1, axes = (1,2)),
                               np.rot90(data, k=2, axes = (1,2)),
                               np.rot90(data, k=3, axes = (1,2))], axis = 0)
        return data

    def augment_data(self):
        # Expand train data
        self.train_data["image"] = self.expand_data(self.train_data["image"])
        self.train_data["edges"] = self.expand_data(self.train_data["edges"])
        # Expand validate data
        self.validate_data["image"] = self.expand_data(self.validate_data["image"])
        self.validate_data["edges"] = self.expand_data(self.validate_data["edges"])

    def shuffle_train_data(self):
        p = np.random.permutation(self.train_data["image"].shape[0])
        self.train_data["image"] = self.train_data["image"][p]
        self.train_data["edges"] = self.train_data["edges"][p]

    def train_cnt(self):
        return self.train_data["image"].shape[0]

    def validate_cnt(self):
        return self.validate_data["image"].shape[0]

    def get_batch(self, index, batch_size, is_training):
        batch_start = index * batch_size
        batch_end = batch_start + batch_size
        batch = {"image" : None, "edges" : None}
        if is_training:
            batch["image"] = self.train_data["image"][batch_start:batch_end, ...]
            batch["edges"] = self.train_data["edges"][batch_start:batch_end, ...]
        else:
            batch["image"] = self.validate_data["image"][batch_start:batch_end, ...]
            batch["edges"] = self.validate_data["edges"][batch_start:batch_end, ...]
        return batch

    def print_info(self):
        print("Dataset type: edge detection")
        print("Source:")
        print("    Images: {}".format(self.data_source["image"]))
        print("    Edges: {}".format(self.data_source["edges"]))

        print("Image width: {}".format(self.width))
        print("Image height: {}".format(self.height))

        print("Train samples: {}".format(self.train_cnt()))
        print("    Edge pixels count: {}".format(np.sum(self.train_data["edges"] > 0)))
        print("Validate samples: {}".format(self.validate_cnt()))
        print("    Edge pixels count: {}".format(np.sum(self.validate_data["edges"] > 0)))

