import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
from functools import partial
import logging
import os

try:
    from . import initializer
    from .utils import load_state
except: 
    import initializer
    from utils import load_state

__all__ = ['ResNeXt', 'resnet50', 'resnet101']


def conv3x3x3(in_planes, out_planes, stride=1):
    # 3x3x3 convolution with padding
    return nn.Conv3d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=False)


def conv1x1x1(in_planes, out_planes, stride=1):
    return nn.Conv3d(in_planes,
                     out_planes,
                     kernel_size=1,
                     stride=stride,
                     bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, downsample=None):
        super().__init__()

        self.conv1 = conv3x3x3(in_planes, planes, stride)
        self.bn1 = nn.BatchNorm3d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3x3(planes, planes)
        self.bn2 = nn.BatchNorm3d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_planes, planes, stride=1, downsample=None):
        super().__init__()

        self.conv1 = conv1x1x1(in_planes, planes)
        self.bn1 = nn.BatchNorm3d(planes)
        self.conv2 = conv3x3x3(planes, planes, stride)
        self.bn2 = nn.BatchNorm3d(planes)
        self.conv3 = conv1x1x1(planes, planes * self.expansion)
        self.bn3 = nn.BatchNorm3d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class ResNet(nn.Module):

    def __init__(self,
                 block,
                 layers,
                 block_inplanes=[64, 128, 256, 512],
                 n_input_channels=3,
                 conv1_t_size=7,
                 conv1_t_stride=1,
                 no_max_pool=False,
                 shortcut_type='B',
                 widen_factor=1.0,
                 num_classes=400, 
                 pretrained=True):
        super().__init__()

        block_inplanes = [int(x * widen_factor) for x in block_inplanes]

        self.in_planes = block_inplanes[0]
        self.no_max_pool = no_max_pool

        self.conv1 = nn.Conv3d(n_input_channels,
                               self.in_planes,
                               kernel_size=(conv1_t_size, 7, 7),
                               stride=(conv1_t_stride, 2, 2),
                               padding=(conv1_t_size // 2, 3, 3),
                               bias=False)
        self.bn1 = nn.BatchNorm3d(self.in_planes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, block_inplanes[0], layers[0],
                                       shortcut_type)
        self.layer2 = self._make_layer(block,
                                       block_inplanes[1],
                                       layers[1],
                                       shortcut_type,
                                       stride=2)
        self.layer3 = self._make_layer(block,
                                       block_inplanes[2],
                                       layers[2],
                                       shortcut_type,
                                       stride=2)
        self.layer4 = self._make_layer(block,
                                       block_inplanes[3],
                                       layers[3],
                                       shortcut_type,
                                       stride=2)

        self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.fc = nn.Linear(block_inplanes[3] * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                nn.init.kaiming_normal_(m.weight,
                                        mode='fan_out',
                                        nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm3d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

         # Initialization

        initializer.xavier(net=self)

        if pretrained:
            pretrained_model=os.path.join(os.path.dirname(os.path.realpath(__file__)), 'pretrained/resnet-101-kinetics.pth')
            logging.info("Network:: graph initialized, loading pretrained model: `{}'".format(pretrained_model))
            assert os.path.exists(pretrained_model), "cannot locate: `{}'".format(pretrained_model)
            pretrained = torch.load(pretrained_model)
            load_state(self, pretrained['state_dict'])
        else:
            logging.info("Network:: graph initialized, use random inilization!")

    def _downsample_basic_block(self, x, planes, stride):
        out = F.avg_pool3d(x, kernel_size=1, stride=stride)
        zero_pads = torch.zeros(out.size(0), planes - out.size(1), out.size(2),
                                out.size(3), out.size(4))
        if isinstance(out.data, torch.cuda.FloatTensor):
            zero_pads = zero_pads.cuda()

        out = torch.cat([out.data, zero_pads], dim=1)

        return out

    def _make_layer(self, block, planes, blocks, shortcut_type, stride=1):
        downsample = None
        if stride != 1 or self.in_planes != planes * block.expansion:
            if shortcut_type == 'A':
                downsample = partial(self._downsample_basic_block,
                                     planes=planes * block.expansion,
                                     stride=stride)
            else:
                downsample = nn.Sequential(
                    conv1x1x1(self.in_planes, planes * block.expansion, stride),
                    nn.BatchNorm3d(planes * block.expansion))

        layers = []
        layers.append(
            block(in_planes=self.in_planes,
                  planes=planes,
                  stride=stride,
                  downsample=downsample))
        self.in_planes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.in_planes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        if not self.no_max_pool:
            x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)

        x = x.view(x.size(0), -1)
        x = self.fc(x)

        return x


def get_fine_tuning_parameters(model, ft_begin_index):
    if ft_begin_index == 0:
        return model.parameters()

    ft_module_names = []
    for i in range(ft_begin_index, 5):
        ft_module_names.append('layer{}'.format(i))
    ft_module_names.append('fc')

    parameters = []
    for k, v in model.named_parameters():
        for ft_module in ft_module_names:
            if ft_module in k:
                parameters.append({'params': v})
                break
        else:
            parameters.append({'params': v, 'lr': 0.0})

    return parameters

def RESNET101(**kwargs):
    """Constructs a ResNet-50 model.
    """
    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
    return model

if __name__ == "__main__":
    import torch
    logging.getLogger().setLevel(logging.DEBUG)
    # ---------
    net1 = RESNET101(num_classes=11, pretrained=True)
    data = torch.randn(1,3,16,224,224)
    output1 = net1(data)
    print (output1.shape)

[((22, 11, 28, 19), 'torch.nn.Conv3d', 'nn.Conv3d', (), '', True, 'import torch.nn as nn\n'), ((32, 11, 36, 32), 'torch.nn.Conv3d', 'nn.Conv3d', (), '', True, 'import torch.nn as nn\n'), ((270, 11, 270, 38), 'torch.randn', 'torch.randn', ({(270, 23, 270, 24): '1', (270, 25, 270, 26): '3', (270, 27, 270, 29): '16', (270, 30, 270, 33): '224', (270, 34, 270, 37): '224'}, {}), '(1, 3, 16, 224, 224)', False, 'import torch\n'), ((46, 19, 46, 41), 'torch.nn.BatchNorm3d', 'nn.BatchNorm3d', ({(46, 34, 46, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((47, 20, 47, 41), 'torch.nn.ReLU', 'nn.ReLU', (), '', True, 'import torch.nn as nn\n'), ((49, 19, 49, 41), 'torch.nn.BatchNorm3d', 'nn.BatchNorm3d', ({(49, 34, 49, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((79, 19, 79, 41), 'torch.nn.BatchNorm3d', 'nn.BatchNorm3d', ({(79, 34, 79, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((81, 19, 81, 41), 'torch.nn.BatchNorm3d', 'nn.BatchNorm3d', ({(81, 34, 81, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((83, 19, 83, 58), 'torch.nn.BatchNorm3d', 'nn.BatchNorm3d', ({(83, 34, 83, 57): 'planes * self.expansion'}, {}), '(planes * self.expansion)', True, 'import torch.nn as nn\n'), ((84, 20, 84, 41), 'torch.nn.ReLU', 'nn.ReLU', (), '', True, 'import torch.nn as nn\n'), ((132, 21, 137, 42), 'torch.nn.Conv3d', 'nn.Conv3d', (), '', True, 'import torch.nn as nn\n'), ((138, 19, 138, 49), 'torch.nn.BatchNorm3d', 'nn.BatchNorm3d', ({(138, 34, 138, 48): 'self.in_planes'}, {}), '(self.in_planes)', True, 'import torch.nn as nn\n'), ((139, 20, 139, 41), 'torch.nn.ReLU', 'nn.ReLU', (), '', True, 'import torch.nn as nn\n'), ((140, 23, 140, 71), 'torch.nn.MaxPool3d', 'nn.MaxPool3d', (), '', True, 'import torch.nn as nn\n'), ((159, 23, 159, 54), 'torch.nn.AdaptiveAvgPool3d', 'nn.AdaptiveAvgPool3d', ({(159, 44, 159, 53): '(1, 1, 1)'}, {}), '((1, 1, 1))', True, 'import torch.nn as nn\n'), ((160, 18, 160, 77), 'torch.nn.Linear', 'nn.Linear', ({(160, 28, 160, 63): 'block_inplanes[3] * block.expansion', (160, 65, 160, 76): 'num_classes'}, {}), '(block_inplanes[3] * block.expansion, num_classes)', True, 'import torch.nn as nn\n'), ((173, 8, 173, 36), 'initializer.xavier', 'initializer.xavier', (), '', False, 'import initializer\n'), ((185, 14, 185, 59), 'torch.nn.functional.avg_pool3d', 'F.avg_pool3d', (), '', True, 'import torch.nn.functional as F\n'), ((191, 14, 191, 53), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((217, 15, 217, 37), 'torch.nn.Sequential', 'nn.Sequential', ({(217, 29, 217, 36): '*layers'}, {}), '(*layers)', True, 'import torch.nn as nn\n'), ((178, 19, 178, 51), 'os.path.exists', 'os.path.exists', ({(178, 34, 178, 50): 'pretrained_model'}, {}), '(pretrained_model)', False, 'import os\n'), ((179, 25, 179, 53), 'torch.load', 'torch.load', ({(179, 36, 179, 52): 'pretrained_model'}, {}), '(pretrained_model)', False, 'import torch\n'), ((180, 12, 180, 54), 'utils.load_state', 'load_state', ({(180, 23, 180, 27): 'self', (180, 29, 180, 53): "pretrained['state_dict']"}, {}), "(self, pretrained['state_dict'])", False, 'from utils import load_state\n'), ((182, 12, 182, 80), 'logging.info', 'logging.info', ({(182, 25, 182, 79): '"""Network:: graph initialized, use random inilization!"""'}, {}), "('Network:: graph initialized, use random inilization!')", False, 'import logging\n'), ((267, 4, 267, 23), 'logging.getLogger', 'logging.getLogger', ({}, {}), '()', False, 'import logging\n'), ((164, 16, 166, 60), 'torch.nn.init.kaiming_normal_', 'nn.init.kaiming_normal_', (), '', True, 'import torch.nn as nn\n'), ((199, 29, 201, 51), 'functools.partial', 'partial', (), '', False, 'from functools import partial\n'), ((168, 16, 168, 46), 'torch.nn.init.constant_', 'nn.init.constant_', ({(168, 34, 168, 42): 'm.weight', (168, 44, 168, 45): '(1)'}, {}), '(m.weight, 1)', True, 'import torch.nn as nn\n'), ((169, 16, 169, 44), 'torch.nn.init.constant_', 'nn.init.constant_', ({(169, 34, 169, 40): 'm.bias', (169, 42, 169, 43): '(0)'}, {}), '(m.bias, 0)', True, 'import torch.nn as nn\n'), ((176, 58, 176, 84), 'os.path.realpath', 'os.path.realpath', ({(176, 75, 176, 83): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((205, 20, 205, 60), 'torch.nn.BatchNorm3d', 'nn.BatchNorm3d', ({(205, 35, 205, 59): 'planes * block.expansion'}, {}), '(planes * block.expansion)', True, 'import torch.nn as nn\n')]