# Path: datasets/DataKeys.py
# SEGMENTATION_LABELS = "segmentation_labels"
# SEGMENTATION_LABELS_ORIGINAL_SIZE = "segmentation_labels_original_size"
# BBOX_GUIDANCE = "bbox_guidance"
# UNSIGNED_DISTANCE_TRANSFORM_GUIDANCE = "unsigned_distance_transform_guidance"
# SIGNED_DISTANCE_TRANSFORM_GUIDANCE = "signed_distance_transform_guidance"
# LASER_GUIDANCE = "laser_guidance"
# IMAGES = "images"
# RAW_IMAGES = "raw_images"
# INPUTS = "inputs"
# IMAGE_FILENAMES = "image_filenames"
# CLASSES = "classes"
# IDS = "ids"  # bounding box number 1,2,3...  or 0 if no bounding box
# DT_METHOD = "dt_method"
# RAW_SEGMENTATION_LABELS = "raw_segmentation_labels"
# CLICK_MAPS = "clicks_maps"
# NEG_CLICKS = "neg_clicks"
# POS_CLICKS = "pos_clicks"
# OBJ_TAGS = "obj_tags"
# FEATUREMAP_LABELS = "featuremap_labels"
# FEATUREMAP_BOXES = "featuremap_boxes"
# IMAGE_ID = "image_id"
# IS_CROWD = "is_crowd"
# SEGMENTATION_MASK = "segmentation_mask"
# SKIP_EXAMPLE = "skip_example"  # If true, filter the corresponding example from the datastream. Useful, when an example
# RAW_IMAGE_SIZES = "raw_image_sizes"  # h, w
# TIME_OFFSETS = "time_offsets"
# SPACE_OFFSETS = "space_offsets"
# SEGMENTATION_INSTANCE_LABELS = "segmentation_instance_labels"
# N_VALID_IDS = "n_valid_ids"
# 
# Path: core/Log.py
# class Stream:
# class Log(object):
#   def __init__(self, log, lvl):
#   def write(self, msg):
#   def flush(self):
#   def initialize(self, logs=[], verbosity=[], formatter=[]):
#   def write(self, msg):
# 
# Path: datasets/util/Util.py
# def flip_coords_horizontal_y0x0y1x1(tensor, original_width):
#   original_width_f = tf.cast(original_width, tensor.dtype)
#   return tf.stack([tensor[..., 0], original_width_f - tensor[..., 3], tensor[..., 2],
#                    original_width_f - tensor[..., 1]], axis=-1)
# 
# def flip_coords_horizontal_x0y0x1y1(tensor, original_width):
#   original_width_f = tf.cast(original_width, tensor.dtype)
#   return tf.stack([original_width_f - tensor[..., 2], tensor[..., 1], original_width_f - tensor[..., 0],
#                    tensor[..., 3]], axis=-1)


import tensorflow as tf
import math
from datasets import DataKeys
from core.Log import log
from datasets.util.Util import flip_coords_horizontal_y0x0y1x1, flip_coords_horizontal_x0y0x1y1


class Augmentor:
  def apply_before_resize(self, tensors):
    return tensors

  def apply_after_resize(self, tensors):
    return tensors

  def batch_apply_before_resize(self, tensors_batch):
    return tensors_batch

  def batch_apply_after_resize(self, tensors_batch):
    return tensors_batch


class GammaAugmentor(Augmentor):
  def __init__(self, gamma_range=(-0.1, 0.1)):
    self.gamma_range = gamma_range

  def apply_after_resize(self, tensors, factor=None):
    """
    Augments the images. Expects it to be in the [0, 1] range
    """
    with tf.name_scope('gamma_augmentor'):
      img = tensors[DataKeys.IMAGES]

      # Sample a gamma factor
      if factor is None:
        factor = self._sample_factor()

# Path: datasets/DataKeys.py
# SEGMENTATION_LABELS = "segmentation_labels"
# SEGMENTATION_LABELS_ORIGINAL_SIZE = "segmentation_labels_original_size"
# BBOX_GUIDANCE = "bbox_guidance"
# UNSIGNED_DISTANCE_TRANSFORM_GUIDANCE = "unsigned_distance_transform_guidance"
# SIGNED_DISTANCE_TRANSFORM_GUIDANCE = "signed_distance_transform_guidance"
# LASER_GUIDANCE = "laser_guidance"
# IMAGES = "images"
# RAW_IMAGES = "raw_images"
# INPUTS = "inputs"
# IMAGE_FILENAMES = "image_filenames"
# CLASSES = "classes"
# IDS = "ids"  # bounding box number 1,2,3...  or 0 if no bounding box
# DT_METHOD = "dt_method"
# RAW_SEGMENTATION_LABELS = "raw_segmentation_labels"
# CLICK_MAPS = "clicks_maps"
# NEG_CLICKS = "neg_clicks"
# POS_CLICKS = "pos_clicks"
# OBJ_TAGS = "obj_tags"
# FEATUREMAP_LABELS = "featuremap_labels"
# FEATUREMAP_BOXES = "featuremap_boxes"
# IMAGE_ID = "image_id"
# IS_CROWD = "is_crowd"
# SEGMENTATION_MASK = "segmentation_mask"
# SKIP_EXAMPLE = "skip_example"  # If true, filter the corresponding example from the datastream. Useful, when an example
# RAW_IMAGE_SIZES = "raw_image_sizes"  # h, w
# TIME_OFFSETS = "time_offsets"
# SPACE_OFFSETS = "space_offsets"
# SEGMENTATION_INSTANCE_LABELS = "segmentation_instance_labels"
# N_VALID_IDS = "n_valid_ids"
# 
# Path: core/Log.py
# class Stream:
# class Log(object):
#   def __init__(self, log, lvl):
#   def write(self, msg):
#   def flush(self):
#   def initialize(self, logs=[], verbosity=[], formatter=[]):
#   def write(self, msg):
# 
# Path: datasets/util/Util.py
# def flip_coords_horizontal_y0x0y1x1(tensor, original_width):
#   original_width_f = tf.cast(original_width, tensor.dtype)
#   return tf.stack([tensor[..., 0], original_width_f - tensor[..., 3], tensor[..., 2],
#                    original_width_f - tensor[..., 1]], axis=-1)
# 
# def flip_coords_horizontal_x0y0x1y1(tensor, original_width):
#   original_width_f = tf.cast(original_width, tensor.dtype)
#   return tf.stack([original_width_f - tensor[..., 2], tensor[..., 1], original_width_f - tensor[..., 0],
#                    tensor[..., 3]], axis=-1)


# Path: datasets/Augmentors.py
import tensorflow as tf
import math
from datasets import DataKeys
from core.Log import log
from datasets.util.Util import flip_coords_horizontal_y0x0y1x1, flip_coords_horizontal_x0y0x1y1



class Augmentor:
  def apply_before_resize(self, tensors):
    return tensors

  def apply_after_resize(self, tensors):
    return tensors

  def batch_apply_before_resize(self, tensors_batch):
    return tensors_batch

  def batch_apply_after_resize(self, tensors_batch):
    return tensors_batch


class GammaAugmentor(Augmentor):
  def __init__(self, gamma_range=(-0.1, 0.1)):
    self.gamma_range = gamma_range

  def apply_after_resize(self, tensors, factor=None):
    """
    Augments the images. Expects it to be in the [0, 1] range
    """
    with tf.name_scope('gamma_augmentor'):
      img = tensors[DataKeys.IMAGES]

      # Sample a gamma factor
      if factor is None:
        factor = self._sample_factor()

      gamma = tf.log(0.5 + 1 / math.sqrt(2) * factor) / tf.log(0.5 - 1 / math.sqrt(2) * factor)