import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo
import torch.nn.functional as F
import torch
import numpy as np
import cv2
import pdb
def sigmoid(x):
    return 1 / (1 + math.exp(-x))


def norm_angle(angle):
    norm_angle = sigmoid(10 * (abs(angle) / 0.7853975 - 1))
    return norm_angle


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


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU()
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(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, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU()
        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 = out + residual
        out = self.relu(out)

        return out

###''' self-attention; relation-attention '''

class ResNet_AT(nn.Module):
    def __init__(self, block, layers, num_classes=1000, end2end=True, at_type=''):
        self.inplanes = 64
        self.end2end = end2end
        super(ResNet_AT, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(0.5)
        self.dropout2 = nn.Dropout(0.6)
        self.alpha = nn.Sequential(nn.Linear(512, 1),
                                   nn.Sigmoid())

        self.beta = nn.Sequential(nn.Linear(1024, 1),
                                  nn.Sigmoid())

        self.pred_fc1 = nn.Linear(512, 7)
        self.pred_fc2 = nn.Linear(1024, 7)
        self.at_type = at_type

        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))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

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

        return nn.Sequential(*layers)

    def forward(self, x='', phrase='train', AT_level='first_level',vectors='',vm='',alphas_from1='',index_matrix=''):

        vs = []
        alphas = []

        assert phrase == 'train' or phrase == 'eval'
        assert AT_level == 'first_level' or AT_level == 'second_level' or AT_level == 'pred'
        if phrase == 'train':
            num_pair = 3

            for i in range(num_pair):
                f = x[:, :, :, :, i]  # x[128,3,224,224]

                f = self.conv1(f)
                f = self.bn1(f)
                f = self.relu(f)
                f = self.maxpool(f)

                f = self.layer1(f)
                f = self.layer2(f)
                f = self.layer3(f)
                f = self.layer4(f)
                f = self.avgpool(f)

                f = f.squeeze(3).squeeze(2)  # f[1, 512, 1, 1] ---> f[1, 512]

                # MN_MODEL(first Level)
                vs.append(f)
                alphas.append(self.alpha(self.dropout(f)))

            vs_stack = torch.stack(vs, dim=2)
            alphas_stack = torch.stack(alphas, dim=2)

            if self.at_type == 'self-attention':
                vm1 = vs_stack.mul(alphas_stack).sum(2).div(alphas_stack.sum(2))
            if self.at_type == 'self_relation-attention':
                vm1 = vs_stack.mul(alphas_stack).sum(2).div(alphas_stack.sum(2))
                betas = []
                for i in range(len(vs)):
                    vs[i] = torch.cat([vs[i], vm1], dim=1)
                    betas.append(self.beta(self.dropout(vs[i])))

                cascadeVs_stack = torch.stack(vs, dim=2)
                betas_stack = torch.stack(betas, dim=2)
                output = cascadeVs_stack.mul(betas_stack * alphas_stack).sum(2).div((betas_stack * alphas_stack).sum(2))

            if self.at_type == 'self-attention':
                vm1 = self.dropout(vm1)
                pred_score = self.pred_fc1(vm1)

            if self.at_type == 'self_relation-attention':
                output = self.dropout2(output)
                pred_score = self.pred_fc2(output)

            return pred_score

        if phrase == 'eval':
            if AT_level == 'first_level':
                f = self.conv1(x)
                f = self.bn1(f)
                f = self.relu(f)
                f = self.maxpool(f)

                f = self.layer1(f)
                f = self.layer2(f)
                f = self.layer3(f)
                f = self.layer4(f)
                f = self.avgpool(f)

                f = f.squeeze(3).squeeze(2)  # f[1, 512, 1, 1] ---> f[1, 512]
                # MN_MODEL(first Level)
                alphas = self.alpha(self.dropout(f))

                return f, alphas

            if AT_level == 'second_level':
                assert self.at_type == 'self_relation-attention'
                vms = index_matrix.permute(1, 0).mm(vm)  # [381, 21783] -> [21783,381] * [381,512] --> [21783, 512]
                vs_cate = torch.cat([vectors, vms], dim=1)

                betas = self.beta(self.dropout(vs_cate))
                ''' keywords: mean_fc ; weight_sourcefc; sum_alpha; weightmean_sourcefc '''
                ''' alpha * beta '''
                weight_catefc = vs_cate.mul(alphas_from1)  # [21570,512] * [21570,1] --->[21570,512]
                alpha_beta = alphas_from1.mul(betas)
                sum_alphabetas = index_matrix.mm(alpha_beta)  # [380,21570] * [21570,1] -> [380,1]
                weightmean_catefc = index_matrix.mm(weight_catefc).div(sum_alphabetas)

                weightmean_catefc = self.dropout2(weightmean_catefc)
                pred_score = self.pred_fc2(weightmean_catefc)

                return pred_score

            if AT_level == 'pred':
                if self.at_type == 'self-attention':
                    pred_score = self.pred_fc1(self.dropout(vm))

                return pred_score

''' self-attention; relation-attention '''
def resnet18_at(pretrained=False, **kwargs):
    # Constructs base a ResNet-18 model.
    model = ResNet_AT(BasicBlock, [2, 2, 2, 2], **kwargs)
    return model

[((20, 11, 21, 43), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((30, 19, 30, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(30, 34, 30, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((31, 20, 31, 29), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((33, 19, 33, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(33, 34, 33, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((61, 21, 61, 75), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((62, 19, 62, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(62, 34, 62, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((63, 21, 64, 53), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((65, 19, 65, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(65, 34, 65, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((66, 21, 66, 77), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((67, 19, 67, 45), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(67, 34, 67, 44): 'planes * 4'}, {}), '(planes * 4)', True, 'import torch.nn as nn\n'), ((68, 20, 68, 29), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((101, 21, 102, 42), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((103, 19, 103, 37), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(103, 34, 103, 36): '64'}, {}), '(64)', True, 'import torch.nn as nn\n'), ((104, 20, 104, 29), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((105, 23, 105, 71), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((110, 23, 110, 46), 'torch.nn.AdaptiveAvgPool2d', 'nn.AdaptiveAvgPool2d', ({(110, 44, 110, 45): '1'}, {}), '(1)', True, 'import torch.nn as nn\n'), ((111, 23, 111, 38), 'torch.nn.Dropout', 'nn.Dropout', ({(111, 34, 111, 37): '0.5'}, {}), '(0.5)', True, 'import torch.nn as nn\n'), ((112, 24, 112, 39), 'torch.nn.Dropout', 'nn.Dropout', ({(112, 35, 112, 38): '0.6'}, {}), '(0.6)', True, 'import torch.nn as nn\n'), ((119, 24, 119, 41), 'torch.nn.Linear', 'nn.Linear', ({(119, 34, 119, 37): '512', (119, 39, 119, 40): '7'}, {}), '(512, 7)', True, 'import torch.nn as nn\n'), ((120, 24, 120, 42), 'torch.nn.Linear', 'nn.Linear', ({(120, 34, 120, 38): '1024', (120, 40, 120, 41): '7'}, {}), '(1024, 7)', True, 'import torch.nn as nn\n'), ((146, 15, 146, 37), 'torch.nn.Sequential', 'nn.Sequential', ({(146, 29, 146, 36): '*layers'}, {}), '(*layers)', True, 'import torch.nn as nn\n'), ((10, 20, 10, 32), 'math.exp', 'math.exp', ({(10, 29, 10, 31): '(-x)'}, {}), '(-x)', False, 'import math\n'), ((113, 35, 113, 52), 'torch.nn.Linear', 'nn.Linear', ({(113, 45, 113, 48): '512', (113, 50, 113, 51): '1'}, {}), '(512, 1)', True, 'import torch.nn as nn\n'), ((114, 35, 114, 47), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((116, 34, 116, 52), 'torch.nn.Linear', 'nn.Linear', ({(116, 44, 116, 48): '1024', (116, 50, 116, 51): '1'}, {}), '(1024, 1)', True, 'import torch.nn as nn\n'), ((117, 34, 117, 46), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((178, 23, 178, 45), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((179, 27, 179, 53), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((135, 16, 136, 67), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((137, 16, 137, 56), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(137, 31, 137, 55): 'planes * block.expansion'}, {}), '(planes * block.expansion)', True, 'import torch.nn as nn\n'), ((190, 34, 190, 56), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((191, 30, 191, 55), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((226, 26, 226, 58), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((126, 41, 126, 58), 'math.sqrt', 'math.sqrt', ({(126, 51, 126, 57): '(2.0 / n)'}, {}), '(2.0 / n)', False, 'import math\n'), ((187, 28, 187, 58), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n')]