# pylint: disable=abstract-method
# import os

import torch
import torch.nn as nn
import torchvision.models as models

# from lib import evoLVe_model


class AdversarialLoss(nn.Module):
  r"""
    Adversarial loss
    https://arxiv.org/abs/1711.10337
    """

  def __init__(self, types='nsgan', target_real_label=1.0, target_fake_label=0.0):
    r"""
        type = nsgan | lsgan | hinge
        """
    super(AdversarialLoss, self).__init__()

    self.type = types
    self.register_buffer('real_label', torch.tensor(target_real_label))
    self.register_buffer('fake_label', torch.tensor(target_fake_label))

    if self.type == 'nsgan':
      self.criterion = nn.BCELoss()

    elif self.type == 'lsgan':
      self.criterion = nn.MSELoss()

    elif self.type == 'hinge':
      self.criterion = nn.ReLU()

  def __call__(self, outputs, is_real, is_disc=None):
    if self.type == 'hinge':
      if is_disc:
        if is_real:
          outputs = -outputs
        return self.criterion(1 + outputs).mean()
      else:
        return (-outputs).mean()

    else:
      labels = (self.real_label if is_real else self.fake_label).expand_as(outputs)
      loss = self.criterion(outputs, labels)
      return loss


class FacialLoss(nn.Module):
  r"""
    Perceptual loss, VGG-based
    https://arxiv.org/abs/1603.08155
    https://github.com/dxyang/StyleTransfer/blob/master/utils.py
    """

  def __init__(self):
    super(FacialLoss, self).__init__()
    self.add_module('evolve', Evolve())
    self.criterion = torch.nn.CosineSimilarity()

  def __call__(self, x, y):
    # Compute features
    x_feat, y_feat = self.evolve(x), self.evolve(y)

    # Compute loss
    return torch.mean(self.criterion(x_feat, y_feat))


class StyleLoss(nn.Module):
  r"""
    Perceptual loss, VGG-based
    https://arxiv.org/abs/1603.08155
    https://github.com/dxyang/StyleTransfer/blob/master/utils.py
    """

  def __init__(self):
    super(StyleLoss, self).__init__()
    self.add_module('vgg', VGG19())
    self.criterion = torch.nn.L1Loss()

  def compute_gram(self, x):
    b, ch, h, w = x.size()
    f = x.view(b, ch, w * h)
    f_T = f.transpose(1, 2)
    G = f.bmm(f_T) / (h * w * ch)

    return G

  def __call__(self, x, y):
    # Compute features
    x_vgg, y_vgg = self.vgg(x), self.vgg(y)

    # Compute loss
    style_loss = 0.0
    style_loss += self.criterion(self.compute_gram(x_vgg['relu2_2']),
                                 self.compute_gram(y_vgg['relu2_2']))
    style_loss += self.criterion(self.compute_gram(x_vgg['relu3_4']),
                                 self.compute_gram(y_vgg['relu3_4']))
    style_loss += self.criterion(self.compute_gram(x_vgg['relu4_4']),
                                 self.compute_gram(y_vgg['relu4_4']))
    style_loss += self.criterion(self.compute_gram(x_vgg['relu5_2']),
                                 self.compute_gram(y_vgg['relu5_2']))

    return style_loss


class PerceptualLoss(nn.Module):
  r"""
    Perceptual loss, VGG-based
    https://arxiv.org/abs/1603.08155
    https://github.com/dxyang/StyleTransfer/blob/master/utils.py
    """

  def __init__(self, weights=[1.0, 1.0, 1.0, 1.0, 1.0]):
    super(PerceptualLoss, self).__init__()
    self.add_module('vgg', VGG19())
    self.criterion = torch.nn.L1Loss()
    self.weights = weights

  def __call__(self, x, y):
    # Compute features
    x_vgg, y_vgg = self.vgg(x), self.vgg(y)

    content_loss = 0.0
    content_loss += self.weights[0] * self.criterion(x_vgg['relu1_1'], y_vgg['relu1_1'])
    content_loss += self.weights[1] * self.criterion(x_vgg['relu2_1'], y_vgg['relu2_1'])
    content_loss += self.weights[2] * self.criterion(x_vgg['relu3_1'], y_vgg['relu3_1'])
    content_loss += self.weights[3] * self.criterion(x_vgg['relu4_1'], y_vgg['relu4_1'])
    content_loss += self.weights[4] * self.criterion(x_vgg['relu5_1'], y_vgg['relu5_1'])

    return content_loss


# class Evolve(torch.nn.Module):

#   def __init__(self):
#     super(Evolve, self).__init__()
#     model_path = 'data/models/backbone_ir50_ms1m_epoch120.pth'
#     if not os.path.isfile(model_path):
#       model_path = os.path.join('/data/uv_inpainting', model_path)

#     self.backbone = evoLVe_model.IR_50([112, 112])
#     self.backbone.eval()
#     checkpoint = torch.load(model_path)
#     self.backbone.load_state_dict(checkpoint)

#     for param in self.parameters():
#       param.requires_grad = False

#   def forward(self, x):
#     # the inputs shoud be in [-1, 1]
#     x = nn.functional.interpolate(x, size=112)
#     return self.backbone(x)


class VGG19(torch.nn.Module):

  def __init__(self):
    super(VGG19, self).__init__()
    features = models.vgg19(pretrained=True).features
    self.relu1_1 = torch.nn.Sequential()
    self.relu1_2 = torch.nn.Sequential()

    self.relu2_1 = torch.nn.Sequential()
    self.relu2_2 = torch.nn.Sequential()

    self.relu3_1 = torch.nn.Sequential()
    self.relu3_2 = torch.nn.Sequential()
    self.relu3_3 = torch.nn.Sequential()
    self.relu3_4 = torch.nn.Sequential()

    self.relu4_1 = torch.nn.Sequential()
    self.relu4_2 = torch.nn.Sequential()
    self.relu4_3 = torch.nn.Sequential()
    self.relu4_4 = torch.nn.Sequential()

    self.relu5_1 = torch.nn.Sequential()
    self.relu5_2 = torch.nn.Sequential()
    self.relu5_3 = torch.nn.Sequential()
    self.relu5_4 = torch.nn.Sequential()

    for x in range(2):
      self.relu1_1.add_module(str(x), features[x])

    for x in range(2, 4):
      self.relu1_2.add_module(str(x), features[x])

    for x in range(4, 7):
      self.relu2_1.add_module(str(x), features[x])

    for x in range(7, 9):
      self.relu2_2.add_module(str(x), features[x])

    for x in range(9, 12):
      self.relu3_1.add_module(str(x), features[x])

    for x in range(12, 14):
      self.relu3_2.add_module(str(x), features[x])

    for x in range(14, 16):
      self.relu3_3.add_module(str(x), features[x])

    for x in range(16, 18):
      self.relu3_4.add_module(str(x), features[x])

    for x in range(18, 21):
      self.relu4_1.add_module(str(x), features[x])

    for x in range(21, 23):
      self.relu4_2.add_module(str(x), features[x])

    for x in range(23, 25):
      self.relu4_3.add_module(str(x), features[x])

    for x in range(25, 27):
      self.relu4_4.add_module(str(x), features[x])

    for x in range(27, 30):
      self.relu5_1.add_module(str(x), features[x])

    for x in range(30, 32):
      self.relu5_2.add_module(str(x), features[x])

    for x in range(32, 34):
      self.relu5_3.add_module(str(x), features[x])

    for x in range(34, 36):
      self.relu5_4.add_module(str(x), features[x])

    # don't need the gradients, just want the features
    for param in self.parameters():
      param.requires_grad = False

  def forward(self, x):
    relu1_1 = self.relu1_1(x)
    relu1_2 = self.relu1_2(relu1_1)

    relu2_1 = self.relu2_1(relu1_2)
    relu2_2 = self.relu2_2(relu2_1)

    relu3_1 = self.relu3_1(relu2_2)
    relu3_2 = self.relu3_2(relu3_1)
    relu3_3 = self.relu3_3(relu3_2)
    relu3_4 = self.relu3_4(relu3_3)

    relu4_1 = self.relu4_1(relu3_4)
    relu4_2 = self.relu4_2(relu4_1)
    relu4_3 = self.relu4_3(relu4_2)
    relu4_4 = self.relu4_4(relu4_3)

    relu5_1 = self.relu5_1(relu4_4)
    relu5_2 = self.relu5_2(relu5_1)
    relu5_3 = self.relu5_3(relu5_2)
    relu5_4 = self.relu5_4(relu5_3)

    out = {
        'relu1_1': relu1_1,
        'relu1_2': relu1_2,
        'relu2_1': relu2_1,
        'relu2_2': relu2_2,
        'relu3_1': relu3_1,
        'relu3_2': relu3_2,
        'relu3_3': relu3_3,
        'relu3_4': relu3_4,
        'relu4_1': relu4_1,
        'relu4_2': relu4_2,
        'relu4_3': relu4_3,
        'relu4_4': relu4_4,
        'relu5_1': relu5_1,
        'relu5_2': relu5_2,
        'relu5_3': relu5_3,
        'relu5_4': relu5_4,
    }
    return out