# Copyright 2022 The Authors
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     https://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import torchvision.transforms as transforms
from PIL import Image


def svhn_transform():
  transform_train = transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
  ])

  transform_test = transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
  ])
  return transform_train, transform_test


def cifar_transform():
  transform_train = transforms.Compose([
      transforms.RandomHorizontalFlip(),
      transforms.RandomCrop(size=[32, 32], padding=4),
      transforms.ToTensor(),
      #transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
  ])

  transform_test = transforms.Compose([
      transforms.ToTensor(),
      #transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
  ])
  return transform_train, transform_test


imagenet_pca = {
    'eigval':
      np.asarray([0.2175, 0.0188, 0.0045]),
    'eigvec':
      np.asarray([
          [-0.5675, 0.7192, 0.4009],
          [-0.5808, -0.0045, -0.8140],
          [-0.5836, -0.6948, 0.4203],
      ])
}


class Lighting(object):
  def __init__(self,
      alphastd,
      eigval=imagenet_pca['eigval'],
      eigvec=imagenet_pca['eigvec']):
    self.alphastd = alphastd
    assert eigval.shape == (3, )
    assert eigvec.shape == (3, 3)
    self.eigval = eigval
    self.eigvec = eigvec

  def __call__(self, img):
    if self.alphastd == 0.:
      return img
    rnd = np.random.randn(3) * self.alphastd
    rnd = rnd.astype('float32')
    v = rnd
    old_dtype = np.asarray(img).dtype
    v = v * self.eigval
    v = v.reshape((3, 1))
    inc = np.dot(self.eigvec, v).reshape((3, ))
    img = np.add(img, inc)
    if old_dtype == np.uint8:
      img = np.clip(img, 0, 255)
    img = Image.fromarray(img.astype(old_dtype), 'RGB')
    return img

  def __repr__(self):
    return self.__class__.__name__ + '()'


def imagenet_transform(img_size=224):
  normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                   [0.229, 0.224, 0.225])
  jitter_param = 0.4
  lighting_param = 0.1

  transform_train = transforms.Compose([
      transforms.RandomResizedCrop(img_size, scale=(0.25, 1)),
      transforms.RandomHorizontalFlip(),
      transforms.ColorJitter(brightness=jitter_param,
                             contrast=jitter_param,
                             saturation=jitter_param),
      Lighting(lighting_param),
      transforms.ToTensor(), normalize
  ])
  transform_test = transforms.Compose([
      transforms.Resize(256),
      transforms.CenterCrop(img_size),
      transforms.ToTensor(), normalize
  ])
  return transform_train, transform_test