# Copyright 2020 The KaiJIN 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.
# ==============================================================================
'''ResNet in PyTorch.

For Pre-activation ResNet, see 'preact_resnet.py'.

Reference:
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
    Deep Residual Learning for Image Recognition. arXiv:1512.03385
'''
import torch
import torch.nn as nn
import torch.nn.functional as F

r"""
  ResNet-20      0.27    8.75
  ResNet-32      0.46    7.51
  ResNet-44      0.66    7.17
  ResNet-56      0.85    6.97
  ResNet-110     1.70    6.43
  ResNet-1202    19.4    7.93
"""

__all__ = ['resnet20', 'resnet32', 'resnet44', 'resnet56', 'resnet110', 'resnet1202']


def conv3x3(in_planes, out_planes, stride=1):
  "3x3 convolution with padding"
  return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                   padding=1, bias=False)


class BasicBlock(nn.Module):
  expansion = 1

  def __init__(self, inplanes, planes, stride=1, downsample=None):
    super(BasicBlock, self).__init__()
    self.conv_1 = conv3x3(inplanes, planes, stride)
    self.bn_1 = nn.BatchNorm2d(planes)
    self.relu = nn.ReLU(inplace=True)
    self.conv_2 = conv3x3(planes, planes)
    self.bn_2 = nn.BatchNorm2d(planes)
    self.downsample = downsample
    self.stride = stride

  def forward(self, x):
    residual = x

    out = self.conv_1(x)
    out = self.bn_1(out)
    out = self.relu(out)

    out = self.conv_2(out)
    out = self.bn_2(out)

    if self.downsample is not None:
      residual = self.downsample(x)

    out += residual
    out = self.relu(out)

    return out


class Bottleneck(nn.Module):
  expansion = 4

  def __init__(self, inplanes, planes, stride=1, downsample=None):
    super(Bottleneck, self).__init__()
    self.conv_1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
    self.bn_1 = nn.BatchNorm2d(planes)
    self.conv_2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                            padding=1, bias=False)
    self.bn_2 = nn.BatchNorm2d(planes)
    self.conv_3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
    self.bn_3 = nn.BatchNorm2d(planes * 4)
    self.relu = nn.ReLU(inplace=True)
    self.downsample = downsample
    self.stride = stride

  def forward(self, x):
    residual = x

    out = self.conv_1(x)
    out = self.bn_1(out)
    out = self.relu(out)

    out = self.conv_2(out)
    out = self.bn_2(out)
    out = self.relu(out)

    out = self.conv_3(out)
    out = self.bn_3(out)

    if self.downsample is not None:
      residual = self.downsample(x)

    out += residual
    out = self.relu(out)

    return out


class ResNet(nn.Module):

  def __init__(self, depth, num_classes, block_name='BasicBlock'):
    super(ResNet, self).__init__()
    # Model type specifies number of layers for CIFAR-10 model
    if block_name == 'BasicBlock':
      assert (depth - 2) % 6 == 0, 'depth should be 6n+2, e.g. 20, 32, 44, 56, 110, 1202'  # nopep8
      n = (depth - 2) // 6
      block = BasicBlock
    elif block_name == 'Bottleneck':
      assert (depth - 2) % 9 == 0, 'depth should be 9n+2, e.g. 20, 29, 47, 56, 110, 1199'  # nopep8
      n = (depth - 2) // 9
      block = Bottleneck
    else:
      raise ValueError('block_name shoule be Basicblock or Bottleneck')

    self.inplanes = 16
    self.conv_1 = nn.Conv2d(3, 16, kernel_size=3, padding=1, bias=False)  # nopep8
    self.bn_1 = nn.BatchNorm2d(16)
    self.relu = nn.ReLU(inplace=True)
    self.stage_1 = self._make_layer(block, 16, n)
    self.stage_2 = self._make_layer(block, 32, n, stride=2)
    self.stage_3 = self._make_layer(block, 64, n, stride=2)
    self.avgpool = nn.AvgPool2d(8)
    self.fc = nn.Linear(64 * block.expansion, num_classes)

    for m in self.modules():
      if isinstance(m, nn.Conv2d):
        # nn.init.xavier_normal(m.weight.data)
        nn.init.kaiming_normal_(m.weight.data)
      elif isinstance(m, nn.BatchNorm2d):
        m.weight.data.fill_(1)
        m.bias.data.zero_()

  def _make_layer(self, block, planes, blocks, stride=1):
    downsample = None
    if stride != 1 or self.inplanes != planes * block.expansion:
      downsample = nn.Sequential(
          nn.Conv2d(self.inplanes, planes * block.expansion,
                    kernel_size=1, stride=stride, bias=False),
          nn.BatchNorm2d(planes * block.expansion),
      )

    layers = []
    layers.append(block(self.inplanes, planes, stride, downsample))
    self.inplanes = planes * block.expansion
    for i in range(1, blocks):
      layers.append(block(self.inplanes, planes))

    return nn.Sequential(*layers)

  def forward(self, x):
    x = self.conv_1(x)
    x = self.bn_1(x)
    x = self.relu(x)     # 32x32

    x = self.stage_1(x)  # 32x32
    x = self.stage_2(x)  # 16x16
    x = self.stage_3(x)  # 8x8

    x = self.avgpool(x)
    x = x.view(x.size(0), -1)
    x = self.fc(x)

    return x


def resnet20(num_classes):
  return ResNet(depth=20, num_classes=num_classes)


def resnet32(num_classes):
  return ResNet(depth=32, num_classes=num_classes)


def resnet44(num_classes):
  return ResNet(depth=44, num_classes=num_classes)


def resnet56(num_classes):
  return ResNet(depth=56, num_classes=num_classes)


def resnet110(num_classes):
  return ResNet(depth=110, num_classes=num_classes)


def resnet1202(num_classes):
  return ResNet(depth=1202, num_classes=num_classes)
