import os

import tensorflow as tf

import augmentations


class CoronetOrgDataset:
    def __init__(self, image_height=512, image_width=512, max_rotation=50,
                 max_shear=0.25, max_pixel_shift=10, max_pixel_scale_change=0.2, shuffle_buffer=1000):
        # General parameters
        self.image_height = image_height
        self.image_width = image_width
        self.shuffle_buffer = shuffle_buffer

        # Augmentation parameters
        self.max_rotation = max_rotation
        self.max_shear = max_shear
        self.max_pixel_shift = max_pixel_shift
        self.max_pixel_scale_change = max_pixel_scale_change

    def train_dataset(self, train_split_file='train.txt', batch_size=1):
        return self._make_dataset(train_split_file, batch_size, True)

    def validation_dataset(self, val_split_file='val.txt', batch_size=1):
        return self._make_dataset(val_split_file, batch_size, False)

    def _make_dataset(self, split_file, batch_size, is_training):
        files, classes, bboxes = self._get_files(split_file)
        count = len(files)
        print("Original file num "+str(count))
        dataset = tf.data.Dataset.from_tensor_slices((files, classes, bboxes))

        if is_training:
            dataset = dataset.shuffle(buffer_size=self.shuffle_buffer)

        # Create and apply map function
        load_and_process = self._get_load_and_process_fn(is_training)
        dataset = dataset.map(load_and_process)

        # Batch data
        dataset = dataset.batch(batch_size)
        return dataset, count, batch_size

    def _get_load_and_process_fn(self, is_training):
        """Creates map function for TF dataset"""
        def load_and_process(path, label, bbox):
            img_file = tf.io.read_file(path)
            if tf.io.is_jpeg(img_file):
                image = tf.image.decode_jpeg(img_file, channels=1)
            else:
                image = tf.image.decode_png(img_file, channels=1)

            # Apply augmentations and/or crop to bbox
            if is_training:
                image, bbox = self._augment_image_and_bbox(image, bbox)
            else:
                image = tf.image.crop_to_bounding_box(image, bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0])

            # Stack to 3-channel, scale to [0, 1] and resize
            image = tf.image.grayscale_to_rgb(image)
            image = tf.cast(image, tf.float32)
            image = image / 255.0
            image = tf.image.resize(image, [self.image_height, self.image_width])
            label = tf.cast(label, dtype=tf.int32)

            return {'image': image, 'label': label}

        return load_and_process

    def _augment_image_and_bbox(self, image, bbox):
        """Apply augmentations to image and bbox"""
        image_shape = tf.cast(tf.shape(image), tf.float32)
        image_height, image_width = image_shape[0], image_shape[1]
        image = augmentations.random_exterior_exclusion(image)
        bbox = augmentations.random_bbox_jitter(bbox, image_height, image_width, self.max_bbox_jitter)
        image, bbox = augmentations.random_rotation(image, self.max_rotation, bbox)
        image, bbox = augmentations.random_shear(image, self.max_shear, bbox)
        image = tf.image.crop_to_bounding_box(image, bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0])
        image = augmentations.random_shift_and_scale(image, self.max_pixel_shift, self.max_pixel_scale_change)
        image = tf.image.random_flip_left_right(image)
        return image, bbox

    def _get_files(self, split_file):
        files, classes, bboxes = [], [], []
        with open(split_file, 'r') as f:
            for line in f.readlines():
                fname, cls, xmin, ymin, xmax, ymax = line.strip('\n').split()
                files.append(fname)
                classes.append(int(cls))
                bboxes.append([int(xmin), int(ymin), int(xmax), int(ymax)])

        return files, classes, bboxes

import os

import tensorflow as tf

import augmentations


class CoronetOrgDataset:
    def __init__(self, image_height=512, image_width=512, max_rotation=50,
                 max_shear=0.25, max_pixel_shift=10, max_pixel_scale_change=0.2, shuffle_buffer=1000):
        # General parameters
        self.image_height = image_height
        self.image_width = image_width
        self.shuffle_buffer = shuffle_buffer

        # Augmentation parameters
        self.max_rotation = max_rotation
        self.max_shear = max_shear
        self.max_pixel_shift = max_pixel_shift
        self.max_pixel_scale_change = max_pixel_scale_change

    def train_dataset(self, train_split_file='train.txt', batch_size=1):
        return self._make_dataset(train_split_file, batch_size, True)

    def validation_dataset(self, val_split_file='val.txt', batch_size=1):
        return self._make_dataset(val_split_file, batch_size, False)

    def _make_dataset(self, split_file, batch_size, is_training):
        files, classes, bboxes = self._get_files(split_file)
        count = len(files)
        print("Original file num "+str(count))
        dataset = tf.data.Dataset.from_tensor_slices((files, classes, bboxes))

        if is_training:
            dataset = dataset.shuffle(buffer_size=self.shuffle_buffer)

        # Create and apply map function
        load_and_process = self._get_load_and_process_fn(is_training)
        dataset = dataset.map(load_and_process)

        # Batch data
        dataset = dataset.batch(batch_size)
        return dataset, count, batch_size

    def _get_load_and_process_fn(self, is_training):
        """Creates map function for TF dataset"""
        def load_and_process(path, label, bbox):
            img_file = tf.io.read_file(path)
            if tf.io.is_jpeg(img_file):
                image = tf.image.decode_jpeg(img_file, channels=1)
            else:
                image = tf.image.decode_png(img_file, channels=1)

            # Apply augmentations and/or crop to bbox
            if is_training:
                image, bbox = self._augment_image_and_bbox(image, bbox)
            else:
                image = tf.image.crop_to_bounding_box(image, bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0])

            # Stack to 3-channel, scale to [0, 1] and resize
            image = tf.image.grayscale_to_rgb(image)
            image = tf.cast(image, tf.float32)
            image = image / 255.0
            image = tf.image.resize(image, [self.image_height, self.image_width])
            label = tf.cast(label, dtype=tf.int32)

            return {'image': image, 'label': label}

        return load_and_process

    def _augment_image_and_bbox(self, image, bbox):
        """Apply augmentations to image and bbox"""
        image_shape = tf.cast(tf.shape(image), tf.float32)
        image_height, image_width = image_shape[0], image_shape[1]
        image = augmentations.random_exterior_exclusion(image)
        bbox = augmentations.random_bbox_jitter(bbox, image_height, image_width, self.max_bbox_jitter)
        image, bbox = augmentations.random_rotation(image, self.max_rotation, bbox)
        image, bbox = augmentations.random_shear(image, self.max_shear, bbox)
        image = tf.image.crop_to_bounding_box(image, bbox[1], bbox[0], bbox[3] - bbox[1], bbox[2] - bbox[0])
        image = augmentations.random_shift_and_scale(image, self.max_pixel_shift, self.max_pixel_scale_change)
        image = tf.image.random_flip_left_right(image)
        return image, bbox

    def _get_files(self, split_file):
        files, classes, bboxes = [], [], []
        with open(split_file, 'r') as f:
            for line in f.readlines():
                fname, cls, xmin, ymin, xmax, ymax = line.strip('\n').split()
                files.append(fname)
                classes.append(int(cls))
                bboxes.append([int(xmin), int(ymin), int(xmax), int(ymax)])

        return files, classes, bboxes

# General parameters # Augmentation parameters # Create and apply map function # Batch data Creates map function for TF dataset # Apply augmentations and/or crop to bbox # Stack to 3-channel, scale to [0, 1] and resize Apply augmentations to image and bbox