'''
CODE TAKEN FROM https://github.com/akamaster/pytorch_resnet_cifar10/blob/master/resnet.py.
WRITTEN BY YERLAN IDELBAYEV, THE DEFINITIONS REFLECT ACCURATE IMPLEMENTATIONS OF THE RESNET
ARCHITECTURES USED IN THE CIFAR-10 EXPERIMENTS OF THE ORIGINAL RESNET PAPER
(Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. Deep Residual Learning for Image Recognition. arXiv:1512.03385).

CODE MODIFICATION MADE TO ALLOW FOR DIFFERENT FORWARD PASS STRUCTURE FOR DISTIL.
    last:   RETURN LAST-LAYER FC EMBEDDINGS IN ADDITION TO TENSOR
    freeze: DO NOT COMPUTE THE COMPUTATION GRAPH UNTIL THE LAST LAYER IS REACHED

ADDITIONAL PARAMETER ADDED TO ALLOW FOR NUMBER OF CHANNELS WHEN INSTANTIATING MODEL. NOTE, HOWEVER, THAT IF THE PROVIDED
VALUE IS NOT 3, THEN THE NUMBER OF PARAMETERS WILL CHANGE FROM THAT OF THE ORIGINAL RESNET PAPER.

PARAMETER COUNTS FOR 3-CHANNEL RESNET
-------------------------------------
name      | layers | params
ResNet20  |    20  | 0.27M
ResNet32  |    32  | 0.46M
ResNet44  |    44  | 0.66M
ResNet56  |    56  | 0.85M
ResNet110 |   110  |  1.7M
ResNet1202|  1202  | 19.4m

'''

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

from torch.autograd import Variable

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


def weight_reset(m):
    if hasattr(m, 'reset_parameters'):
        m.reset_parameters()


def _weights_init(m):
    classname = m.__class__.__name__
    #print(classname)
    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
        init.kaiming_normal_(m.weight)

class LambdaLayer(nn.Module):
    def __init__(self, lambd):
        super(LambdaLayer, self).__init__()
        self.lambd = lambd

    def forward(self, x):
        return self.lambd(x)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, option='A'):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != planes:
            if option == 'A':
                """
                For CIFAR10 ResNet paper uses option A.
                """
                self.shortcut = LambdaLayer(lambda x:
                                            F.pad(x[:, :, ::2, ::2], (0, 0, 0, 0, planes//4, planes//4), "constant", 0))
            elif option == 'B':
                self.shortcut = nn.Sequential(
                     nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
                     nn.BatchNorm2d(self.expansion * planes)
                )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10, channels=3):
        super(ResNet, self).__init__()
        self.in_planes = 16

        self.conv1 = nn.Conv2d(channels, 16, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)
        self.linear = nn.Linear(64, num_classes)

        self.apply(_weights_init)

    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 forward(self, x, last=False, freeze=False):

        if freeze:
            with torch.no_grad():
                out = F.relu(self.bn1(self.conv1(x)))
                out = self.layer1(out)
                out = self.layer2(out)
                out = self.layer3(out)
                out = F.avg_pool2d(out, out.size()[3])
                last_layer_emb = out.view(out.size(0), -1)
        else:
            out = F.relu(self.bn1(self.conv1(x)))
            out = self.layer1(out)
            out = self.layer2(out)
            out = self.layer3(out)
            out = F.avg_pool2d(out, out.size()[3])
            last_layer_emb = out.view(out.size(0), -1)
        out = self.linear(last_layer_emb)
        
        if last:
            return out, last_layer_emb
        else:
            return out

    def get_embedding_dim(self):
        return 64

    def reset(self):
        self.apply(weight_reset)


def resnet20(num_classes=10, channels=3):
    return ResNet(BasicBlock, [3, 3, 3], num_classes, channels)


def resnet32(num_classes=10, channels=3):
    return ResNet(BasicBlock, [5, 5, 5], num_classes, channels)


def resnet44(num_classes=10, channels=3):
    return ResNet(BasicBlock, [7, 7, 7], num_classes, channels)


def resnet56(num_classes=10, channels=3):
    return ResNet(BasicBlock, [9, 9, 9], num_classes, channels)


def resnet110(num_classes=10, channels=3):
    return ResNet(BasicBlock, [18, 18, 18], num_classes, channels)


def resnet1202(num_classes=10, channels=3):
    return ResNet(BasicBlock, [200, 200, 200], num_classes, channels)


if __name__ == "__main__":
    
    model = resnet20()
    size_inp = 512
    test_inp = torch.rand(1, 3, size_inp, size_inp)
    out = model(test_inp)