# -*- encoding: utf-8 -*-
#Time        :2022/02/24 18:14:15
#Author      :Hao Chen
#FileName    :trans_lib.py
#Version     :2.0

import cv2
import torch
import numpy as np
import albumentations as A


def get_albu_transforms(type="train", img_size = (192, 192)):
    if type == 'train':
        compose = [
            # A.VerticalFlip(p=0.5),
            # A.HorizontalFlip(p=0.5),

            A.ShiftScaleRotate(shift_limit=0.2, scale_limit=(-0.2, 0.2),
                                rotate_limit=5, p=0.5),

            A.OneOf([
                A.GridDistortion(num_steps=1, distort_limit=0.3, p=1.0),
                A.ElasticTransform(alpha=2, sigma=5, p=1.0)
            ], p=0.5),

            A.Resize(img_size[0], img_size[1])]
    else:
        compose = [A.Resize(img_size[0], img_size[1])]

    return A.Compose(compose, p=1.0, additional_targets={'image2': 'image',
                                                         'image3': 'image',
                                                         'image4': 'image',
                                                         'image5': 'image',
                                                         'image6': 'image',
                                                         "mask2": "mask"})




# Beta function
def gamma_concern(img, gamma):
    mean = torch.mean(img)

    img = (img - mean) * gamma
    img = img + mean
    img = torch.clip(img, 0, 1)

    return img

def gamma_power(img, gamma, direction=0):
    if direction == 1:
        img = 1 - img
    img = torch.pow(img, gamma)

    img = img / torch.max(img)
    if direction == 1:
        img = 1 - img

    return img

def gamma_exp(img, gamma, direction=0):
    if direction == 1:
        img = 1 - img

    img = torch.exp(img * gamma)
    img = img / torch.max(img)

    if direction == 1:
        img = 1 - img
    return img