# 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
    return [self.apply_after_resize(x, factor) for x in tensors_batch]


class FlipAugmentor(Augmentor):
  def __init__(self, p=0.5):
    """
    :param p: The probability that the image will be flipped.
    """
    self.p = p

  def apply_after_resize(self, tensors, doit=None):
    with tf.name_scope("flip_augmentor"):
      aug_tensors = tensors.copy()

      if doit is None:
        doit = self._sample_doit()

      def maybe_flip(key_, image_flip, y0x0y1x1=False):
        if key_ in tensors:
          val = tensors[key_]
          if image_flip:
            flipped = tf.image.flip_left_right(val)
          elif y0x0y1x1:
            flipped = flip_coords_horizontal_y0x0y1x1(val, tf.shape(tensors[DataKeys.IMAGES])[1])
          else:
            flipped = flip_coords_horizontal_x0y0x1y1(val, tf.shape(tensors[DataKeys.IMAGES])[1])
          aug = tf.cond(doit, lambda: flipped, lambda: val)
          aug_tensors[key_] = aug

      keys_to_flip = [DataKeys.IMAGES, DataKeys.SEGMENTATION_LABELS, DataKeys.BBOX_GUIDANCE,

# 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

    return [self.apply_after_resize(x, factor) for x in tensors_batch]


class FlipAugmentor(Augmentor):
  def __init__(self, p=0.5):
    """
    :param p: The probability that the image will be flipped.
    """
    self.p = p

  def apply_after_resize(self, tensors, doit=None):
    with tf.name_scope("flip_augmentor"):
      aug_tensors = tensors.copy()

      if doit is None:
        doit = self._sample_doit()

      def maybe_flip(key_, image_flip, y0x0y1x1=False):
        if key_ in tensors:
          val = tensors[key_]
          if image_flip:
            flipped = tf.image.flip_left_right(val)
          elif y0x0y1x1:
            flipped = flip_coords_horizontal_y0x0y1x1(val, tf.shape(tensors[DataKeys.IMAGES])[1])
          else:
            flipped = flip_coords_horizontal_x0y0x1y1(val, tf.shape(tensors[DataKeys.IMAGES])[1])
          aug = tf.cond(doit, lambda: flipped, lambda: val)
          aug_tensors[key_] = aug

      keys_to_flip = [DataKeys.IMAGES, DataKeys.SEGMENTATION_LABELS, DataKeys.BBOX_GUIDANCE,

                      DataKeys.SEGMENTATION_LABELS_ORIGINAL_SIZE, DataKeys.LASER_GUIDANCE, DataKeys.SEGMENTATION_MASK,