import torch
import torch.nn as nn
import torch.nn.functional as F

from os.path import join as pjoin
from collections import OrderedDict


def weight_standardize(w, dim, eps):
    """Subtracts mean and divides by standard deviation."""
    w = w - torch.mean(w, dim=dim)
    w = w / (torch.std(w, dim=dim) + eps)
    return w


def np2th(weights, conv=False):
    """Possibly convert HWIO to OIHW."""
    if conv:
        weights = weights.transpose([3, 2, 0, 1])
    return torch.from_numpy(weights)


class StdConv2d(nn.Conv2d):
    def forward(self, x):
        w = weight_standardize(self.weight, [0, 1, 2], 1e-5)
        return F.conv2d(x, w, self.bias, self.stride, self.padding,
                        self.dilation, self.groups)


def conv3x3(in_channels, out_channels, stride=1, groups=1, bias=False):
    return StdConv2d(in_channels,
                     out_channels,
                     kernel_size=3,
                     stride=stride,
                     padding=1,
                     bias=bias,
                     groups=groups)


def conv1x1(in_channels, out_channels, stride=1, bias=False):
    return StdConv2d(in_channels,
                     out_channels,
                     kernel_size=1,
                     stride=stride,
                     padding=0,
                     bias=bias)


class PreActBottleneck(nn.Module):
    """Pre-activation (v2) bottleneck block.
    """
    def __init__(self,
                 in_channels,
                 out_channels=None,
                 mid_channels=None,
                 stride=1):
        super().__init__()
        out_channels = out_channels or in_channels
        mid_channels = mid_channels or out_channels // 4

        self.gn1 = nn.GroupNorm(32, mid_channels, eps=1e-6)
        self.conv1 = conv1x1(in_channels, mid_channels, bias=False)
        self.gn2 = nn.GroupNorm(32, mid_channels, eps=1e-6)
        self.conv2 = conv3x3(mid_channels, mid_channels, stride,
                             bias=False)  # Original code has it on conv1!!
        self.gn3 = nn.GroupNorm(32, out_channels, eps=1e-6)
        self.conv3 = conv1x1(mid_channels, out_channels, bias=False)
        self.relu = nn.ReLU(inplace=True)

        if (stride != 1 or in_channels != out_channels):
            # Projection also with pre-activation according to paper.
            self.downsample = conv1x1(in_channels,
                                      out_channels,
                                      stride,
                                      bias=False)
            self.gn_proj = nn.GroupNorm(out_channels, out_channels)

    def forward(self, x):

        # Residual branch
        residual = x
        if hasattr(self, 'downsample'):
            residual = self.downsample(x)
            residual = self.gn_proj(residual)

        # Unit's branch
        y = self.relu(self.gn1(self.conv1(x)))
        y = self.relu(self.gn2(self.conv2(y)))
        y = self.gn3(self.conv3(y))

        y = self.relu(residual + y)
        return y


class ResNetV2(nn.Module):
    """Implementation of Pre-activation (v2) ResNet mode."""
    def __init__(self, block_units, width_factor):
        super().__init__()
        width = int(64 * width_factor)
        self.width = width
        self.downsample = 16  # four stride=2 conv2d layer

        # The following will be unreadable if we split lines.
        # pylint: disable=line-too-long
        self.root = nn.Sequential(
            OrderedDict([('conv',
                          StdConv2d(3,
                                    width,
                                    kernel_size=7,
                                    stride=2,
                                    bias=False,
                                    padding=3)),
                         ('gn', nn.GroupNorm(32, width, eps=1e-6)),
                         ('relu', nn.ReLU(inplace=True)),
                         ('pool',
                          nn.MaxPool2d(kernel_size=3, stride=2, padding=0))]))

        self.body = nn.Sequential(
            OrderedDict([
                ('block1',
                 nn.Sequential(
                     OrderedDict([('unit1',
                                   PreActBottleneck(in_channels=width,
                                                    out_channels=width * 4,
                                                    mid_channels=width))] +
                                 [(f'unit{i:d}',
                                   PreActBottleneck(in_channels=width * 4,
                                                    out_channels=width * 4,
                                                    mid_channels=width))
                                  for i in range(2, block_units[0] + 1)], ))),
                ('block2',
                 nn.Sequential(
                     OrderedDict([('unit1',
                                   PreActBottleneck(in_channels=width * 4,
                                                    out_channels=width * 8,
                                                    mid_channels=width * 2,
                                                    stride=2))] +
                                 [(f'unit{i:d}',
                                   PreActBottleneck(in_channels=width * 8,
                                                    out_channels=width * 8,
                                                    mid_channels=width * 2))
                                  for i in range(2, block_units[1] + 1)], ))),
                ('block3',
                 nn.Sequential(
                     OrderedDict([('unit1',
                                   PreActBottleneck(in_channels=width * 8,
                                                    out_channels=width * 16,
                                                    mid_channels=width * 4,
                                                    stride=2))] +
                                 [(f'unit{i:d}',
                                   PreActBottleneck(in_channels=width * 16,
                                                    out_channels=width * 16,
                                                    mid_channels=width * 4))
                                  for i in range(2, block_units[2] + 1)], ))),
            ]))

    def forward(self, x):
        x = self.root(x)
        x = self.body(x)
        return x


def resnet50():
    return ResNetV2(block_units=(3, 4, 9), width_factor=1)