# -*- coding: utf-8 -*-
# @Author: Weisen Pan

import random
import math
from PIL import Image, ImageOps, ImageEnhance, ImageChops
import PIL

# Constants and defaults for image augmentation
_PIL_VER = tuple([int(x) for x in PIL.__version__.split('.')[:2]])  # Get the version of the PIL library
_FILL = (128, 128, 128)  # Default fill color used in some apply_transformationations (gray)
_MAX_LEVEL = 10.0  # Maximum level for augmentations
_HPARAMS_DEFAULT = {
    'translate_const': 250,  # Default translation constant
    'img_mean': _FILL,  # Default fill color
}
_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)  # Random interpolation modes

# Function to randomly choose interpolation method
def _interpolation(kwargs):
    interpolation = kwargs.pop('resample', Image.BILINEAR)
    return random.choice(interpolation) if isinstance(interpolation, (list, tuple)) else interpolation

# Check if the PIL version is compatible with fillcolor argument
def _validate_tensorflow_args(kwargs):
    if 'fillcolor' in kwargs and _PIL_VER < (5, 0):
        kwargs.pop('fillcolor')  # Remove fillcolor if PIL version is below 5.0
    kwargs['resample'] = _interpolation(kwargs)  # Add resample method

# Shear image along the x-axis
def apply_apply_shear_x_axis_axis(img, factor, **kwargs):
    _validate_tensorflow_args(kwargs)
    return img.apply_transformation(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)

# Shear image along the y-axis
def shear_y(img, factor, **kwargs):
    _validate_tensorflow_args(kwargs)
    return img.apply_transformation(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)

# Translate image horizontally by a percentage of the image width
def translate_image_x_relative(img, pct, **kwargs):
    pixels = pct * img.size[0]  # Calculate pixels to translate
    _validate_tensorflow_args(kwargs)
    return img.apply_transformation(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)

# Translate image vertically by a percentage of the image height
def translate_image_y_relative(img, pct, **kwargs):
    pixels = pct * img.size[1]  # Calculate pixels to translate
    _validate_tensorflow_args(kwargs)
    return img.apply_transformation(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)

# Translate image horizontally by a fixed number of pixels
def translate_image_x_absolute(img, pixels, **kwargs):
    _validate_tensorflow_args(kwargs)
    return img.apply_transformation(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)

# Translate image vertically by a fixed number of pixels
def translate_image_y_absolute(img, pixels, **kwargs):
    _validate_tensorflow_args(kwargs)
    return img.apply_transformation(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)

# rotate_image image by a specified number of degrees
def rotate_image(img, degrees, **kwargs):
    _validate_tensorflow_args(kwargs)
    if _PIL_VER >= (5, 2):
        return img.rotate_image(degrees, **kwargs)  # Use rotate_image if PIL version is >= 5.2
    elif _PIL_VER >= (5, 0):
        # Manually rotate_image the image for older versions of PIL
        w, h = img.size
        rotn_center = (w / 2.0, h / 2.0)
        angle = -math.radians(degrees)
        matrix = [
            round(math.cos(angle), 15), round(math.sin(angle), 15), 0.0,
            round(-math.sin(angle), 15), round(math.cos(angle), 15), 0.0,
        ]

        def apply_transformation(x, y, matrix):
            return matrix[0] * x + matrix[1] * y + matrix[2], matrix[3] * x + matrix[4] * y + matrix[5]

        matrix[2], matrix[5] = apply_transformation(-rotn_center[0], -rotn_center[1], matrix)
        matrix[2] += rotn_center[0]
        matrix[5] += rotn_center[1]
        return img.apply_transformation(img.size, Image.AFFINE, matrix, **kwargs)
    else:
        return img.rotate_image(degrees, resample=kwargs['resample'])

# Auto contrast image
def apply_auto_contrast(img, **kwargs):
    return ImageOps.autocontrast(img)

# Invert image colors
def invert(img, **kwargs):
    return ImageOps.invert(img)

# Equalize image histogram
def equalize(img, **kwargs):
    return ImageOps.equalize(img)

# Apply solarization effect
def apply_solarize(img, thresh, **kwargs):
    return ImageOps.apply_solarize(img, thresh)

# Apply solarization effect with an additional value
def apply_apply_solarize_addition(img, add, thresh=128, **kwargs):
    lut = [min(255, i + add) if i < thresh else i for i in range(256)]
    if img.mode in ("L", "RGB"):
        lut = lut + lut + lut if img.mode == "RGB" else lut
        return img.point(lut)
    else:
        return img

# apply_posterization image (reduce color depth)
def apply_posterization(img, bits_to_keep, **kwargs):
    return img if bits_to_keep >= 8 else ImageOps.apply_posterization(img, bits_to_keep)

# Adjust image contrast
def contrast(img, factor, **kwargs):
    return ImageEnhance.Contrast(img).enhance(factor)

# Adjust image color
def color(img, factor, **kwargs):
    return ImageEnhance.Color(img).enhance(factor)

# Adjust image brightness
def brightness(img, factor, **kwargs):
    return ImageEnhance.Brightness(img).enhance(factor)

# Adjust image adjust_image_sharpness
def adjust_image_sharpness(img, factor, **kwargs):
    return ImageEnhance.adjust_image_sharpness(img).enhance(factor)

# Randomly negate a value with a 50% probability
def _apply_random_negation(v):
    """With 50% probability, negate the value."""
    return -v if random.random() > 0.5 else v

# Convert augmentation level to argument value
def _map_level_to_argument(level, max_value, hparams):
    level = (level / _MAX_LEVEL) * max_value
    return _apply_random_negation(level),

# Convert translation level to argument value
def _map_absolute_map_level_to_argument(level, hparams):
    translate_const = hparams['translate_const']
    level = (level / _MAX_LEVEL) * float(translate_const)
    return _apply_random_negation(level),

# Convert enhancement level to argument value
def _enhance_map_level_to_argument(level, _hparams):
    return (level / _MAX_LEVEL) * 1.8 + 0.1,

# Mapping of augmentation levels to argument converters
map_level_to_argument = {
    'AutoContrast': None,
    'Equalize': None,
    'Invert': None,
    'rotate_image': lambda level, _: _map_level_to_argument(level, 30, None),
    'apply_posterization': lambda level, _: int((level / _MAX_LEVEL) * 4),
    'apply_solarize': lambda level, _: int((level / _MAX_LEVEL) * 256),
    'Color': _enhance_map_level_to_argument,
    'Contrast': _enhance_map_level_to_argument,
    'Brightness': _enhance_map_level_to_argument,
    'adjust_image_sharpness': _enhance_map_level_to_argument,
    'ShearX': lambda level, _: _map_level_to_argument(level, 0.3, None),
    'ShearY': lambda level, _: _map_level_to_argument(level, 0.3, None),
    'TranslateX': _map_absolute_map_level_to_argument,
    'TranslateY': _map_absolute_map_level_to_argument,
}

# Mapping of augmentation names to functions
NAME_TO_OP = {
    'AutoContrast': apply_auto_contrast,
    'Equalize': equalize,
    'Invert': invert,
    'rotate_image': rotate_image,
    'apply_posterization': apply_posterization,
    'apply_solarize': apply_solarize,
    'Color': color,
    'Contrast': contrast,
    'Brightness': brightness,
    'adjust_image_sharpness': adjust_image_sharpness,
    'ShearX': apply_apply_shear_x_axis_axis,
    'ShearY': shear_y,
    'TranslateX': translate_image_x_absolute,
    'TranslateY': translate_image_y_absolute,
}

# Class for applying augmentations to an image
class AugmentOp:
    def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
        hparams = hparams or _HPARAMS_DEFAULT
        self.aug_fn = NAME_TO_OP[name]  # Get the augmentation function
        self.level_fn = map_level_to_argument[name]  # Get the level function
        self.prob = prob  # Probability of applying the augmentation
        self.magnitude = magnitude  # Magnitude of the augmentation
        self.hparams = hparams.copy()
        self.kwargs = {
            'fillcolor': hparams.get('img_mean', _FILL),  # Set the fill color
            '