# Copyright 2020 Lorna Authors. All Rights Reserved.
# 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
#
#       http://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 torch
import torch.nn as nn

from .utils import get_model_params
from .utils import load_pretrained_weights
from .utils import vgg_params

cfgs = {
  "A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
  "B": [64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"],
  "D": [64, 64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"],
  "E": [64, 64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"],
}


class VGG(nn.Module):

  def __init__(self, global_params=None):
    """ An VGGNet model. Most easily loaded with the .from_name or .from_pretrained methods

    Args:
      global_params (namedtuple): A set of GlobalParams shared between blocks

    Examples:
      model = VGG.from_pretrained('vgg11')
    """

    super(VGG, self).__init__()
    self._global_params = global_params

    cfg = self._global_params.cfg
    batch_norm = self._global_params.batch_norm
    dropout_rate = self._global_params.dropout_rate
    num_classes = self._global_params.num_classes

    self.features = make_layers(cfgs[cfg], batch_norm)
    self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
    self.classifier = nn.Sequential(
      nn.Linear(512 * 7 * 7, 4096),
      nn.ReLU(True),
      nn.Dropout(dropout_rate),
      nn.Linear(4096, 4096),
      nn.ReLU(True),
      nn.Dropout(dropout_rate),
      nn.Linear(4096, num_classes),
    )

    for m in self.modules():
      if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
          nn.init.constant_(m.bias, 0)
      elif isinstance(m, nn.BatchNorm2d):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)
      elif isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, 0, 0.01)
        nn.init.constant_(m.bias, 0)

  def extract_features(self, inputs):
    """ Returns output of the final convolution layer """
    x = self.features(inputs)
    return x

  def forward(self, x):
    x = self.features(x)
    x = self.avgpool(x)
    x = torch.flatten(x, 1)
    x = self.classifier(x)
    return x

  @classmethod
  def from_name(cls, model_name, override_params=None):
    cls._check_model_name_is_valid(model_name)
    global_params = get_model_params(model_name, override_params)
    return cls(global_params)

  @classmethod
  def from_pretrained(cls, model_name, num_classes=1000):
    model = cls.from_name(model_name, override_params={'num_classes': num_classes})
    load_pretrained_weights(model, model_name, load_fc=(num_classes == 1000))
    return model

  @classmethod
  def get_image_size(cls, model_name):
    cls._check_model_name_is_valid(model_name)
    _, res, _ = vgg_params(model_name)
    return res

  @classmethod
  def _check_model_name_is_valid(cls, model_name):
    """ Validates model name. None that pretrained weights are only available for
    the first four models (vgg{i} for i in 11,13,16,19) at the moment. """
    valid_models = ['vgg' + str(i) for i in ["11", "11_bn",
                                             "13", "13_bn",
                                             "16", "16_bn",
                                             "19", "19_bn"]]
    if model_name not in valid_models:
      raise ValueError('model_name should be one of: ' + ', '.join(valid_models))


def make_layers(cfg, batch_norm):
  layers = []
  in_channels = 3
  for v in cfg:
    if v == 'M':
      layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
    else:
      conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
      if batch_norm:
        layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
      else:
        layers += [conv2d, nn.ReLU(inplace=True)]
      in_channels = v
  return nn.Sequential(*layers)
