#!/usr/bin/env python
# coding: utf-8
#
# Author:   Kazuto Nakashima
# URL:      http://kazuto1011.github.io
# Created:  2017-11-19

from __future__ import absolute_import, print_function

from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F
from mlassistant.core import Model

try:
    from encoding.nn import SyncBatchNorm

    _BATCH_NORM = SyncBatchNorm
except:
    _BATCH_NORM = nn.BatchNorm2d

_BOTTLENECK_EXPANSION = 4


class ConvBlock(Model):
    """
    Cascade of 2D convolution, batch norm, and ReLU.
    """

    BATCH_NORM = _BATCH_NORM

    def __init__(self, in_ch, out_ch, kernel_size, stride, padding, dilation, relu=True):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(
            in_ch, out_ch, kernel_size, stride, padding, dilation, bias=False)
        self.bn = _BATCH_NORM(out_ch, eps=1e-5, momentum=1 - 0.999)
        self.relu_activation = nn.ReLU()
        self.relu = relu

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        if self.relu:
            x = self.relu_activation(x)
        return x


class ResnetBlock(nn.Module):
    """
    Bottleneck block of MSRA ResNet.
    """

    def __init__(self, in_ch, out_ch, stride, dilation, downsample):
        super(ResnetBlock, self).__init__()
        mid_ch = out_ch // _BOTTLENECK_EXPANSION
        self.reduce = ConvBlock(in_ch, mid_ch, 1, stride, 0, 1, True)
        self.conv3x3 = ConvBlock(mid_ch, mid_ch, 3, 1, dilation, dilation, True)
        self.increase = ConvBlock(mid_ch, out_ch, 1, 1, 0, 1, False)
        self.shortcut = (
            ConvBlock(in_ch, out_ch, 1, stride, 0, 1, False) if downsample else nn.Identity())

    def forward(self, x):
        h = self.reduce(x)
        h = self.conv3x3(h)
        h = self.increase(h)
        h = h +  self.shortcut(x)
        return F.relu(h)


class ResnetLayer(nn.Module):
    """
    Residual layer with multi grids
    """

    def __init__(self, n_layers, in_ch, out_ch, stride, dilation, multi_grids=None):
        super(ResnetLayer, self).__init__()
        self.layers = nn.ModuleList()
        if multi_grids is None:
            multi_grids = [1 for _ in range(n_layers)]
        else:
            assert n_layers == len(multi_grids)

        for i in range(n_layers):
            self.layers.append(
                ResnetBlock(
                    in_ch=(in_ch if i == 0 else out_ch),
                    out_ch=out_ch,
                    stride=(stride if i == 0 else 1),
                    dilation=dilation * multi_grids[i],
                    downsample=(True if i == 0 else False)))

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


class PoolBlock(nn.Module):
    """
    The 1st conv layer.
    Note that the max pooling is different from both MSRA and FAIR ResNet.
    """

    def __init__(self, out_ch):
        super(PoolBlock, self).__init__()
        self.conv1 = ConvBlock(1, out_ch, 7, 2, 3, 1)
        self.pool = nn.MaxPool2d(3, 2, 1, ceil_mode=True)

    def forward(self, x):
        x = self.conv1(x)
        x = self.pool(x)
        return x


class ResNet(nn.Module):

    def __init__(self, n_classes, n_blocks):
        super(ResNet, self).__init__()
        ch = [64 * 2**p for p in range(6)]
        self.layer1 = PoolBlock(ch[0])
        self.layer2 = ResnetLayer(n_blocks[0], ch[0], ch[2], 1, 1)
        self.layer3 = ResnetLayer(n_blocks[1], ch[2], ch[3], 2, 1)
        self.layer4 = ResnetLayer(n_blocks[2], ch[3], ch[4], 2, 1)
        self.pool5 = nn.AdaptiveAvgPool2d(1)
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(ch[5], n_classes)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.pool5(x)
        x = self.flatten(x)
        x = self.fc(x)
        return x


if __name__ == "__main__":
    model = ResNet(n_classes=1000, n_blocks=[3, 4, 23, 3])
    model.eval()
    image = torch.randn(1, 3, 224, 224)

    print(model)
    print("input:", image.shape)
    print("output:", model(image).shape)