import torch.nn as nn
import torch.nn.functional as F
import torch.utils.model_zoo as model_zoo
# import math
# from torch.nn import init


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


class PreActBlock(nn.Module):
    '''Pre-activation version of the BasicBlock.'''
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, droprate=0):
        super(PreActBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.conv1 = conv3x3(in_planes, planes, stride)
        self.drop = nn.Dropout(p=droprate) if droprate > 0 else None
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes)

        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes,
                          kernel_size=1, stride=stride, bias=False)
            )

    def forward(self, x):
        x = self.bn1(x)
        out = F.relu(x)
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv1(out)
        if self.drop is not None:
            out = self.drop(out)
        out = self.bn2(out)
        out = self.conv2(F.relu(out))
        out += shortcut
        return out

class BNBlock(nn.Module):
    '''Pre-activation version of the BasicBlock.'''

    def __init__(self, in_planes):
        super(BNBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)

    def forward(self, x):
        out = self.bn1(x)
        return out


class PreActBottleneck(nn.Module):
    '''Pre-activation version of the original Bottleneck module.'''
    expansion = 4

    def __init__(self, in_planes, planes , stride=1,  droprate=None):
        super(PreActBottleneck, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
                               stride=stride, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, self.expansion *
                               planes, kernel_size=1, bias=False)

        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes,
                          kernel_size=1, stride=stride, bias=False)
            )

    def forward(self, x):
        out = F.relu(self.bn1(x))
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv1(out)
        out = self.conv2(F.relu(self.bn2(out)))
        out = self.conv3(F.relu(self.bn3(out)))
        out += shortcut
        return out


class PreActResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10, in_channels=3):
        super(PreActResNet, self).__init__()
        self.in_planes = 64
        last_planes = 512*block.expansion

        self.conv1 = conv3x3(in_channels, 64)
        self.stage1 = self._make_layer(block, 64, num_blocks[0], stride=1)
        self.stage2 = self._make_layer(block, 128, num_blocks[1], stride=2)
        self.stage3 = self._make_layer(block, 256, num_blocks[2], stride=2)
        self.stage4 = self._make_layer(block, 512, num_blocks[3], stride=2)
        self.bn_last = nn.BatchNorm2d(last_planes)
        # self.bn_last = self._make_last(last_planes)
        self.fc = nn.Linear(last_planes, num_classes, 1)
        # self.last = nn.Linear(last_planes, num_classes)  # last layer

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def _make_last(self, last_planes):
        layers = []
        layers.append(BNBlock(last_planes))
        return nn.Sequential(*layers)

    def features(self, x):
        out = self.conv1(x)
        out = self.stage1(out)
        out = self.stage2(out)
        out = self.stage3(out)
        out = self.stage4(out)
        return out

    def logits(self, x):
        x = self.fc(x)
        # x = self.last(x)
        return x

    def forward(self, x):
        x = self.features(x)
        x = self.bn_last(x)
        x = F.relu(x)
        x = F.adaptive_avg_pool2d(x, 1)
        x = self.logits(x.view(x.size(0), -1))
        return x


class PreActResNet_cifar(nn.Module):
    def __init__(self, block, num_blocks, filters, num_classes=10, droprate=0):
        super(PreActResNet_cifar, self).__init__()
        self.in_planes = 16
        last_planes = filters[2]*block.expansion

        self.conv1 = conv3x3(3, self.in_planes)
        self.stage1 = self._make_layer(
            block, filters[0], num_blocks[0], stride=1, droprate=droprate)
        self.stage2 = self._make_layer(
            block, filters[1], num_blocks[1], stride=2, droprate=droprate)
        self.stage3 = self._make_layer(
            block, filters[2], num_blocks[2], stride=2, droprate=droprate)
        self.bn_last = nn.BatchNorm2d(last_planes)
        self.fc = nn.Linear(last_planes, num_classes, bias=False)
        # self.last = nn.Linear(last_planes, num_classes, bias=False)

        """
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
                # m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                init.kaiming_normal(m.weight)
                m.bias.data.zero_()
        """

    def _make_layer(self, block, planes, num_blocks, stride, droprate):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride, droprate))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def features(self, x):
        out = self.conv1(x)
        out = self.stage1(out)
        out = self.stage2(out)
        out = self.stage3(out)
        return out

    def logits(self, x):
        x = self.fc(x)
        # x = self.last(x)
        return x

    def forward(self, x):
        out = self.features(x)
        out = F.relu(self.bn_last(out))
        out = F.avg_pool2d(out, 8)
        out = self.logits(out.view(out.size(0), -1))
        return out


def resnet18():
    # model_urls = {
    #     'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    #     'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    #     'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    #     'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    #     'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
    # }
    print('Unpretrained resnet18!')
    model = PreActResNet(PreActBlock, [2, 2, 2, 2], num_classes=10)
    # model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
    return model
