from functools import partial
from collections import defaultdict

import torch
from torch import nn
from torch.nn import functional as F
import torch.utils.model_zoo as model_zoo

__all__ = ['nonlocalresnet3d50']

model_urls = {
    'kinetics-400': defaultdict(lambda: None, {
        'resnet3d18': 'http://pretorched-x.csail.mit.edu/models/resnet3d18_kinetics-b5673e4a.pth',
        'resnet3d34': 'http://pretorched-x.csail.mit.edu/models/resnet3d34_kinetics-133ec9c4.pth',
        'nonlocalresnet3d50': 'http://pretorched-x.csail.mit.edu/models/resnet3d50_kinetics-aad059c9.pth',
        'resnet3d101': 'http://pretorched-x.csail.mit.edu/models/resnet3d101_kinetics-a6ddd22a.pth',
        'resnet3d152': 'http://pretorched-x.csail.mit.edu/models/resnet3d152_kinetics-8ae08d3f.pth',
    }),
    'moments': defaultdict(lambda: None, {
        'resnet3d50': 'http://pretorched-x.csail.mit.edu/models/resnet3d50_16seg_moments-22f4fe61.pth',
    }),
}

num_classes = {'kinetics-400': 400, 'moments': 339}

pretrained_settings = defaultdict(dict)
input_sizes = {}
means = {}
stds = {}

for model_name in __all__:
    input_sizes[model_name] = [3, 224, 224]
    means[model_name] = [0.485, 0.456, 0.406]
    stds[model_name] = [0.229, 0.224, 0.225]

for model_name in __all__:
    if model_name in ['ResNet3D']:
        continue
    for dataset, urls in model_urls.items():
        pretrained_settings[model_name][dataset] = {
            'input_space': 'RGB',
            'input_range': [0, 1],
            'url': urls[model_name],
            'std': stds[model_name],
            'mean': means[model_name],
            'num_classes': num_classes[dataset],
            'input_size': input_sizes[model_name],
        }


class _NonLocalBlockND(nn.Module):
    def __init__(self, in_channels, inter_channels=None, dimension=3, mode='embedded_gaussian',
                 sub_sample=False, bn_layer=True):
        super(_NonLocalBlockND, self).__init__()

        assert dimension in [1, 2, 3]
        assert mode in ['embedded_gaussian', 'gaussian', 'dot_product', 'concatenation']

        # print(f'Dimension: {dimension}, mode: {mode}')

        self.mode = mode
        self.dimension = dimension
        self.sub_sample = sub_sample

        self.in_channels = in_channels
        self.inter_channels = inter_channels

        if self.inter_channels is None:
            self.inter_channels = in_channels // 2
            if self.inter_channels == 0:
                self.inter_channels = 1

        if dimension == 3:
            conv_nd = nn.Conv3d
            max_pool = nn.MaxPool3d
            bn = nn.BatchNorm3d
        elif dimension == 2:
            conv_nd = nn.Conv2d
            max_pool = nn.MaxPool2d
            bn = nn.BatchNorm2d
        else:
            conv_nd = nn.Conv1d
            max_pool = nn.MaxPool1d
            bn = nn.BatchNorm1d

        self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                         kernel_size=1, stride=1, padding=0)

        if bn_layer:
            self.W = nn.Sequential(
                conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
                        kernel_size=1, stride=1, padding=0),
                bn(self.in_channels)
            )
            nn.init.constant_(self.W[1].weight, 0)
            nn.init.constant_(self.W[1].bias, 0)
        else:
            self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
                             kernel_size=1, stride=1, padding=0)
            nn.init.constant_(self.W.weight, 0)
            nn.init.constant_(self.W.bias, 0)

        self.theta = None
        self.phi = None
        self.concat_project = None

        if mode in ['embedded_gaussian', 'dot_product', 'concatenation']:
            self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                                 kernel_size=1, stride=1, padding=0)
            self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                               kernel_size=1, stride=1, padding=0)

            if mode == 'embedded_gaussian':
                self.operation_function = self._embedded_gaussian
            elif mode == 'dot_product':
                self.operation_function = self._dot_product
            elif mode == 'concatenation':
                self.operation_function = self._concatenation
                self.concat_project = nn.Sequential(
                    nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),
                    nn.ReLU()
                )
        elif mode == 'gaussian':
            self.operation_function = self._gaussian

        if sub_sample:
            self.g = nn.Sequential(self.g, max_pool(kernel_size=2))
            if self.phi is None:
                self.phi = max_pool(kernel_size=2)
            else:
                self.phi = nn.Sequential(self.phi, max_pool(kernel_size=2))

    def forward(self, x):
        """
        :param x: (b, c, t, h, w)
        :return:
        """

        func = getattr(self, f'_{self.mode}')
        output = func(x)
        return output

    def _embedded_gaussian(self, x):
        batch_size = x.size(0)

        # g=>(b, c, t, h, w)->(b, 0.5c, t, h, w)->(b, thw, 0.5c)
        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)

        # theta=>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, thw, 0.5c)
        # phi  =>(b, c, t, h, w)[->(b, 0.5c, t, h, w)]->(b, 0.5c, thw)
        # f=>(b, thw, 0.5c)dot(b, 0.5c, twh) = (b, thw, thw)
        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
        theta_x = theta_x.permute(0, 2, 1)
        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
        f = torch.matmul(theta_x, phi_x)
        f_div_C = F.softmax(f, dim=-1)

        # (b, thw, thw)dot(b, thw, 0.5c) = (b, thw, 0.5c)->(b, 0.5c, t, h, w)->(b, c, t, h, w)
        y = torch.matmul(f_div_C, g_x)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)
        z = W_y + x

        return z

    def _gaussian(self, x):
        batch_size = x.size(0)
        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)

        theta_x = x.view(batch_size, self.in_channels, -1)
        theta_x = theta_x.permute(0, 2, 1)

        if self.sub_sample:
            phi_x = self.phi(x).view(batch_size, self.in_channels, -1)
        else:
            phi_x = x.view(batch_size, self.in_channels, -1)

        f = torch.matmul(theta_x, phi_x)
        f_div_C = F.softmax(f, dim=-1)

        y = torch.matmul(f_div_C, g_x)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)
        z = W_y + x

        return z

    def _dot_product(self, x):
        batch_size = x.size(0)

        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)

        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
        theta_x = theta_x.permute(0, 2, 1)
        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)
        f = torch.matmul(theta_x, phi_x)
        N = f.size(-1)
        f_div_C = f / N

        y = torch.matmul(f_div_C, g_x)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)
        z = W_y + x

        return z

    def _concatenation(self, x):
        batch_size = x.size(0)

        g_x = self.g(x).view(batch_size, self.inter_channels, -1)
        g_x = g_x.permute(0, 2, 1)

        # (b, c, N, 1)
        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1, 1)
        # (b, c, 1, N)
        phi_x = self.phi(x).view(batch_size, self.inter_channels, 1, -1)

        h = theta_x.size(2)
        w = phi_x.size(3)
        theta_x = theta_x.repeat(1, 1, 1, w)
        phi_x = phi_x.repeat(1, 1, h, 1)

        concat_feature = torch.cat([theta_x, phi_x], dim=1)
        f = self.concat_project(concat_feature)
        b, _, h, w = f.size()
        f = f.view(b, h, w)

        N = f.size(-1)
        f_div_C = f / N

        y = torch.matmul(f_div_C, g_x)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])
        W_y = self.W(y)
        z = W_y + x

        return z


class NonLocalBlock1D(_NonLocalBlockND):
    def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', sub_sample=False, bn_layer=True):
        super(NonLocalBlock1D, self).__init__(in_channels,
                                              inter_channels=inter_channels,
                                              dimension=1, mode=mode,
                                              sub_sample=sub_sample,
                                              bn_layer=bn_layer)


class NonLocalBlock2D(_NonLocalBlockND):
    def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', sub_sample=False, bn_layer=True):
        super(NonLocalBlock2D, self).__init__(in_channels,
                                              inter_channels=inter_channels,
                                              dimension=2, mode=mode,
                                              sub_sample=sub_sample,
                                              bn_layer=bn_layer)


class NonLocalBlock3D(_NonLocalBlockND):
    def __init__(self, in_channels, inter_channels=None, mode='embedded_gaussian', sub_sample=False, bn_layer=True):
        super(NonLocalBlock3D, self).__init__(in_channels,
                                              inter_channels=inter_channels,
                                              dimension=3, mode=mode,
                                              sub_sample=sub_sample,
                                              bn_layer=bn_layer)


class MNISTNonLocalNet(nn.Module):
    def __init__(self):
        super().__init__()

        self.convs = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),

            NonLocalBlock2D(in_channels=32),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),

            NonLocalBlock2D(in_channels=64),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )

        self.fc = nn.Sequential(
            nn.Linear(in_features=128*3*3, out_features=256),
            nn.ReLU(),
            nn.Dropout(0.5),

            nn.Linear(in_features=256, out_features=10)
        )

    def forward(self, x):
        batch_size = x.size(0)
        output = self.convs(x).view(batch_size, -1)
        output = self.fc(output)
        return output


def conv3x3x3(in_planes, out_planes, stride=1):
    """3x3x3 Factored Spatial-Temporal convolution with padding."""
    return nn.Conv3d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=False)


def downsample_basic_block(x, planes, stride):
    out = F.avg_pool3d(x, kernel_size=1, stride=stride)
    zero_pads = torch.Tensor(
        out.size(0), planes - out.size(1), out.size(2), out.size(3),
        out.size(4)).zero_()
    if isinstance(out.data, torch.cuda.FloatTensor):
        zero_pads = zero_pads.cuda()
        zero_pads.cuda()
    out = torch.cat([out, zero_pads], dim=1)
    # out = Variable(torch.cat([out.data, zero_pads], dim=1))
    return out


class NonLocalBasicBlock(nn.Module):
    expansion = 1
    Conv3d = staticmethod(conv3x3x3)

    def __init__(self, inplanes, planes, stride=1, downsample=None, nonlocal_layer=False):
        super(NonLocalBasicBlock, self).__init__()
        self.conv1 = self.Conv3d(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm3d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = self.Conv3d(planes, planes)
        self.bn2 = nn.BatchNorm3d(planes)

        self.stride = stride
        self.downsample = downsample
        self.nonlocal_layer = nonlocal_layer

        if nonlocal_layer:
            self.nonlocalblock = NonLocalBlock3D(planes)

    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)

        if self.nonlocal_layer:
            out = self.nonlocalblock(out)

        return out


class NonLocalBottleneck(nn.Module):
    expansion = 4
    Conv3d = nn.Conv3d

    def __init__(self, inplanes, planes, stride=1, downsample=None, nonlocal_layer=False):
        super(NonLocalBottleneck, self).__init__()
        self.conv1 = self.Conv3d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm3d(planes)
        self.conv2 = self.Conv3d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm3d(planes)
        self.conv3 = self.Conv3d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm3d(planes * 4)
        self.relu = nn.ReLU(inplace=True)

        self.stride = stride
        self.downsample = downsample
        self.nonlocal_layer = nonlocal_layer

        if nonlocal_layer:
            self.nonlocalblock = NonLocalBlock3D(planes * 4)

    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)

        if self.nonlocal_layer:
            out = self.nonlocalblock(out)

        return out


class NonLocalResNet3D(nn.Module):

    Conv3d = nn.Conv3d

    def __init__(self,
                 block,
                 layers,
                 nonlocal_layers,
                 shortcut_type='A',
                 num_classes=339):
        self.inplanes = 64
        super().__init__()
        self.conv1 = nn.Conv3d(3, 64, kernel_size=7, stride=(1, 2, 2), padding=(3, 3, 3), bias=False)
        self.bn1 = nn.BatchNorm3d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0], nonlocal_layers[0], shortcut_type)
        self.layer2 = self._make_layer(block, 128, layers[1], nonlocal_layers[1], shortcut_type, stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], nonlocal_layers[2], shortcut_type, stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], nonlocal_layers[3], shortcut_type, stride=2)
        self.avgpool = nn.AdaptiveAvgPool3d(1)
        self.last_linear = nn.Linear(512 * block.expansion, num_classes)

        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, self.Conv3d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
            elif isinstance(m, nn.BatchNorm3d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, nonlocal_blocks, shortcut_type, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            if shortcut_type == 'A':
                downsample = partial(
                    downsample_basic_block,
                    planes=planes * block.expansion,
                    stride=stride)
            else:
                downsample = nn.Sequential(
                    self.Conv3d(
                        self.inplanes,
                        planes * block.expansion,
                        kernel_size=1,
                        stride=stride,
                        bias=False),
                    nn.BatchNorm3d(planes * block.expansion))

        nonlocal_freq = blocks // nonlocal_blocks if nonlocal_blocks != 0 else -1

        layers = []
        for i in range(0, blocks):
            layers.append(block(self.inplanes, planes, stride=stride, downsample=downsample,
                                nonlocal_layer=(i % nonlocal_freq == 0 and nonlocal_freq > 0)))
            if i == 0:
                stride = 1
                downsample = None
                self.inplanes = planes * block.expansion

        return nn.Sequential(*layers)

    def features(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

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

    def logits(self, features):
        x = self.avgpool(features)
        x = x.view(x.size(0), -1)
        x = self.last_linear(x)
        return x

    def forward(self, input):
        x = self.features(input)
        x = self.logits(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('last_linear')

    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 nonlocalresnet3d(**kwargs):
    """Constructs a ResNet-18 model.
    """
    model = NonLocalResNet3D(NonLocalBasicBlock, [1, 1, 1, 1], **kwargs)
    return model


def nonlocalresnet3d18(**kwargs):
    """Constructs a NonLocalResNet3D-18 model.
    """
    model = NonLocalResNet3D(NonLocalBasicBlock, [2, 2, 2, 2], **kwargs)
    return model


def nonlocalresnet3d34(**kwargs):
    """Constructs a NonLocalResNet3D-34 model.
    """
    model = NonLocalResNet3D(NonLocalBasicBlock, [3, 4, 6, 3], **kwargs)
    return model


def nonlocalresnet3d50(num_classes=339, num_nonlocal_blocks=5, pretrained='kinetics-400', **kwargs):
    """Constructs a NonLocalResNet3D-50 model.
    """
    if num_nonlocal_blocks == 5:
        nonlocal_blocks = [0, 2, 3, 0]
    elif num_nonlocal_blocks == 10:
        nonlocal_blocks = [0, 4, 6, 0]

    model = NonLocalResNet3D(NonLocalBottleneck, [3, 4, 6, 3], nonlocal_blocks, **kwargs)
    if pretrained is not None:
        settings = pretrained_settings['nonlocalresnet3d50'][pretrained]
        model.load_state_dict(model_zoo.load_url(settings['url']), strict=False)
        model.input_space = settings['input_space']
        model.input_size = settings['input_size']
        model.input_range = settings['input_range']
        model.mean = settings['mean']
        model.std = settings['std']
    return model


def nonlocalresnet3d101(**kwargs):
    """Constructs a NonLocalResNet3D-101 model.
    """
    model = NonLocalResNet3D(NonLocalBottleneck, [3, 4, 23, 3], **kwargs)
    return model


def nonlocalresnet3d152(**kwargs):
    """Constructs a NonLocalResNet3D-101 model.
    """
    model = NonLocalResNet3D(NonLocalBottleneck, [3, 8, 36, 3], **kwargs)
    return model


def nonlocalresnet3d200(**kwargs):
    """Constructs a NonLocalResNet3D-200 model.
    """
    model = NonLocalResNet3D(NonLocalBottleneck, [3, 24, 36, 3], **kwargs)
    return model


if __name__ == '__main__':
    batch_size = 2
    num_frames = 2
    num_classes = 339
    img_feature_dim = 512
    frame_size = 224
    input_var = torch.randn(batch_size, 3, num_frames, 224, 224)

    model = nonlocalresnet3d50(num_classes=num_classes)
    print(model)
    model = torch.nn.DataParallel(model).cuda()
    # print(model)
    # print(model.layer1)
    # print('===========')
    # print(model.layer2)
    # print('===========')
    # print(model.layer3)
    # print('===========')
    # print(model.layer4)

    output = model(input_var)
    print(output.shape)
    model = nonlocalresnet3d50(num_classes=num_classes, num_nonlocal_blocks=10)
    print(model)
    output = model(input_var)
    print(output.shape)
    # model.cpu()
    # del model

    # model = nonlocalresnet3d(num_classes=num_classes).cuda()
    # output = model(input_var)
    # print(output.shape)
    # model.cpu()
    # del model

    # model = nonlocalresnet3d18(num_classes=num_classes).cuda()
    # output = model(input_var)
    # print(output.shape)
    # model.cpu()
    # del model

    # model = nonlocalresnet3d34(num_classes=num_classes).cuda()
    # output = model(input_var)
    # print(output.shape)
    # model.cpu()
    # del model

    # model = nonlocalresnet3d101(num_classes=num_classes).cuda()
    # output = model(input_var)
    # print(output.shape)
    # model.cpu()
    # del model

    # model = nonlocalresnet3d152(num_classes=num_classes).cuda()
    # output = model(input_var)
    # print(output.shape)
    # model.cpu()
    # del model

    # model = nonlocalresnet3d200(num_classes=num_classes).cuda()
    # output = model(input_var)
    # print(output.shape)
    # model.cpu()
    # del model

    # mode_list = ['concatenation', 'embedded_gaussian', 'gaussian', 'dot_product', ]

    # for mode in mode_list:
    #     print(mode)
    #     img = Variable(torch.zeros(2, 4, 5))
    #     net = NONLocalBlock1D(4, mode=mode, sub_sample=True)
    #     out = net(img)
    #     # print(out.size())

    #     img = Variable(torch.zeros(2, 4, 10, 10))
    #     net = NONLocalBlock2D(4, mode=mode, sub_sample=False, bn_layer=False)
    #     out = net(img)
    #     print(out.size())

    #     img = Variable(torch.zeros(2, 4, 5, 4, 5))
    #     net = NONLocalBlock3D(4, mode=mode)
    #     out = net(img)
    #     print(out.size())

    #     img = Variable(torch.randn(3, 1, 28, 28))
    #     net = Network()
    #     out = net(img)
    #     print(out.size())
