'''EfficientNet in PyTorch.

Paper: "EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks".

Reference: https://github.com/keras-team/keras-applications/blob/master/keras_applications/efficientnet.py
'''
import torch
import torch.nn as nn
import torch.nn.functional as F


def swish(x):
  return x * x.sigmoid()


def drop_connect(x, drop_ratio):
  keep_ratio = 1.0 - drop_ratio
  mask = torch.empty([x.shape[0], 1, 1, 1], dtype=x.dtype, device=x.device)
  mask.bernoulli_(keep_ratio)
  x.div_(keep_ratio)
  x.mul_(mask)
  return x


class SE(nn.Module):
  '''Squeeze-and-Excitation block with Swish.'''

  def __init__(self, in_channels, se_channels):
    super(SE, self).__init__()
    self.se1 = nn.Conv2d(in_channels, se_channels,
                         kernel_size=1, bias=True)
    self.se2 = nn.Conv2d(se_channels, in_channels,
                         kernel_size=1, bias=True)

  def forward(self, x):
    out = F.adaptive_avg_pool2d(x, (1, 1))
    out = swish(self.se1(out))
    out = self.se2(out).sigmoid()
    out = x * out
    return out


class Block(nn.Module):
  '''expansion + depthwise + pointwise + squeeze-excitation'''

  def __init__(self,
               in_channels,
               out_channels,
               kernel_size,
               stride,
               expand_ratio=1,
               se_ratio=0.,
               drop_rate=0.):
    super(Block, self).__init__()
    self.stride = stride
    self.drop_rate = drop_rate
    self.expand_ratio = expand_ratio

    # Expansion
    channels = expand_ratio * in_channels
    self.conv1 = nn.Conv2d(in_channels,
                           channels,
                           kernel_size=1,
                           stride=1,
                           padding=0,
                           bias=False)
    self.bn1 = nn.BatchNorm2d(channels)

    # Depthwise conv
    self.conv2 = nn.Conv2d(channels,
                           channels,
                           kernel_size=kernel_size,
                           stride=stride,
                           padding=(1 if kernel_size == 3 else 2),
                           groups=channels,
                           bias=False)
    self.bn2 = nn.BatchNorm2d(channels)

    # SE layers
    se_channels = int(in_channels * se_ratio)
    self.se = SE(channels, se_channels)

    # Output
    self.conv3 = nn.Conv2d(channels,
                           out_channels,
                           kernel_size=1,
                           stride=1,
                           padding=0,
                           bias=False)
    self.bn3 = nn.BatchNorm2d(out_channels)

    # Skip connection if in and out shapes are the same (MV-V2 style)
    self.has_skip = (stride == 1) and (in_channels == out_channels)

  def forward(self, x):
    out = x if self.expand_ratio == 1 else swish(self.bn1(self.conv1(x)))
    out = swish(self.bn2(self.conv2(out)))
    out = self.se(out)
    out = self.bn3(self.conv3(out))
    if self.has_skip:
      if self.training and self.drop_rate > 0:
        out = drop_connect(out, self.drop_rate)
      out = out + x
    return out


class EfficientNet(nn.Module):
  def __init__(self, cfg, num_classes=10):
    super(EfficientNet, self).__init__()
    self.cfg = cfg
    self.conv1 = nn.Conv2d(3,
                           32,
                           kernel_size=3,
                           stride=1,
                           padding=1,
                           bias=False)
    self.bn1 = nn.BatchNorm2d(32)
    self.layers = self._make_layers(in_channels=32)
    self.linear = nn.Linear(cfg['out_channels'][-1], num_classes)

  def _make_layers(self, in_channels):
    layers = []
    cfg = [self.cfg[k] for k in ['expansion', 'out_channels', 'num_blocks', 'kernel_size',
                                 'stride']]
    b = 0
    blocks = sum(self.cfg['num_blocks'])
    for expansion, out_channels, num_blocks, kernel_size, stride in zip(*cfg):
      strides = [stride] + [1] * (num_blocks - 1)
      for stride in strides:
        drop_rate = self.cfg['drop_connect_rate'] * b / blocks
        layers.append(
            Block(in_channels,
                  out_channels,
                  kernel_size,
                  stride,
                  expansion,
                  se_ratio=0.25,
                  drop_rate=drop_rate))
        in_channels = out_channels
    return nn.Sequential(*layers)

  def forward(self, x):
    out = swish(self.bn1(self.conv1(x)))
    out = self.layers(out)
    out = F.adaptive_avg_pool2d(out, 1)
    out = out.view(out.size(0), -1)
    dropout_rate = self.cfg['dropout_rate']
    if self.training and dropout_rate > 0:
      out = F.dropout(out, p=dropout_rate)
    out = self.linear(out)
    return out


def efficientnet_b0():
  cfg = {
      'num_blocks': [1, 2, 2, 3, 3, 4, 1],
      'expansion': [1, 6, 6, 6, 6, 6, 6],
      'out_channels': [16, 24, 40, 80, 112, 192, 320],
      'kernel_size': [3, 3, 5, 3, 5, 5, 3],
      'stride': [1, 2, 2, 2, 1, 2, 1],
      'dropout_rate': 0.2,
      'drop_connect_rate': 0.2,
  }
  return EfficientNet(cfg)


def test():
  net = efficientnet_b0()
  x = torch.randn(2, 3, 32, 32)
  y = net(x)
  print(y.shape)


if __name__ == '__main__':
  test()
