# print("网络模型输入尺寸448*448\n"
#       "网络结构采用Exploring Self-attention for Image Recognition\n"
#       "自注意层采用Stand-Alone Self-Attention in Vision Models标量编码\n"
#       "位置编码采用On the Relationship between Self-Attention and Convolutional Layers相对向量位置编码\n"
#       "good luck\n"
#       )
# print("san_patchwise用于跨媒体检索")
# 18879270
# x1 torch.Size([1, 3, 448, 448])
# x2 torch.Size([1, 64, 448, 448])
# x3 torch.Size([1, 64, 224, 224])
# x4 torch.Size([1, 256, 112, 112])
# x5 torch.Size([1, 512, 56, 56])
# x6 torch.Size([1, 1024, 28, 28])
# x7 torch.Size([1, 2048, 14, 14])
# avgpool torch.Size([1, 2048, 1, 1])
# view torch.Size([1, 2048])
# fc torch.Size([1, 200])

import torch
import torch.nn as nn

from modules import Subtraction, Subtraction2, Aggregation


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


def position(H, W, is_cuda=True):
    if is_cuda:
        loc_w = torch.linspace(-1.0, 1.0, W).cuda().unsqueeze(0).repeat(H, 1)
        loc_h = torch.linspace(-1.0, 1.0, H).cuda().unsqueeze(1).repeat(1, W)
    else:
        loc_w = torch.linspace(-1.0, 1.0, W).unsqueeze(0).repeat(H, 1)
        loc_h = torch.linspace(-1.0, 1.0, H).unsqueeze(1).repeat(1, W)
    loc = torch.cat([loc_w.unsqueeze(0), loc_h.unsqueeze(0)], 0).unsqueeze(0)
    return loc


class SAM(nn.Module):
    def __init__(self, sa_type, in_planes, rel_planes, out_planes, share_planes, kernel_size=3, stride=1, dilation=1):
        super(SAM, self).__init__()
        self.sa_type, self.kernel_size, self.stride = sa_type, kernel_size, stride
        # self.in_planes, self.rel_planes, self.out_planes, self.share_planes = in_planes, rel_planes, out_planes, share_planes
        self.conv1 = nn.Conv2d(in_planes, rel_planes, kernel_size=1)
        self.conv2 = nn.Conv2d(in_planes, rel_planes, kernel_size=1)
        self.conv3 = nn.Conv2d(in_planes, out_planes, kernel_size=1)
        # 定义自更新m,   m*a=1, m.size() = x1.size()
        self.m = nn.Parameter(torch.randn(1,kernel_size * kernel_size * rel_planes,1,1).cuda())
        if sa_type == 0:
            self.conv_w = nn.Sequential(nn.BatchNorm2d(rel_planes + 2), nn.ReLU(inplace=True),
                                        nn.Conv2d(rel_planes + 2, rel_planes, kernel_size=1, bias=False),
                                        nn.BatchNorm2d(rel_planes), nn.ReLU(inplace=True),
                                        nn.Conv2d(rel_planes, out_planes // share_planes, kernel_size=1))
            self.conv_p = nn.Conv2d(2, 2, kernel_size=1)
            self.subtraction = Subtraction(kernel_size, stride, (dilation * (kernel_size - 1) + 1) // 2, dilation, pad_mode=1)
            self.subtraction2 = Subtraction2(kernel_size, stride, (dilation * (kernel_size - 1) + 1) // 2, dilation, pad_mode=1)
            self.softmax = nn.Softmax(dim=-2)
        else:
            self.conv_w = nn.Sequential(nn.BatchNorm2d(rel_planes * (pow(kernel_size, 2) + 1)), nn.ReLU(inplace=True),
                                        nn.Conv2d(rel_planes * (pow(kernel_size, 2) + 1), out_planes // share_planes, kernel_size=1, bias=False),
                                        nn.BatchNorm2d(out_planes // share_planes), nn.ReLU(inplace=True),
                                        nn.Conv2d(out_planes // share_planes, pow(kernel_size, 2) * out_planes // share_planes, kernel_size=1))

            # self.conv_w1 = nn.Sequential(nn.BatchNorm2d(rel_planes * (pow(kernel_size, 2))), nn.ReLU(inplace=True),
            #                             nn.Conv2d(rel_planes * (pow(kernel_size, 2)), out_planes // share_planes,
            #                                       kernel_size=1, bias=False),
            #                             nn.BatchNorm2d(out_planes // share_planes), nn.ReLU(inplace=True),
            #                             nn.Conv2d(out_planes // share_planes,
            #                                       pow(kernel_size, 2) * out_planes // share_planes, kernel_size=1))

            self.unfold_i = nn.Unfold(kernel_size=1, dilation=dilation, padding=0, stride=stride)
            self.unfold_j = nn.Unfold(kernel_size=kernel_size, dilation=dilation, padding=0, stride=stride)
            self.pad = nn.ReflectionPad2d(kernel_size // 2)
        self.aggregation = Aggregation(kernel_size, stride, (dilation * (kernel_size - 1) + 1) // 2, dilation, pad_mode=1)

    def forward(self, x):
        # print("x",x.size())
        # print(self.sa_type, self.in_planes, self.rel_planes, self.out_planes, self.share_planes,self.kernel_size, self.stride)
        x1, x2, x3 = self.conv1(x), self.conv2(x), self.conv3(x)
        # print("x1",x1.size())
        # print("x2", x2.size())
        # print("x3", x3.size())

        if self.sa_type == 0:  # pairwise
            p = self.conv_p(position(x.shape[2], x.shape[3], x.is_cuda))
            w = self.softmax(self.conv_w(torch.cat([self.subtraction2(x1, x2), self.subtraction(p).repeat(x.shape[0], 1, 1, 1)], 1)))
            # aggregation: x3 torch.Size([1, 16, 224, 224])
            # aggregation: w  torch.Size([1, 2, 9, 50176])
        else:  # patchwise
            if self.stride != 1:
                x1 = self.unfold_i(x1)
            x1 = x1.view(x.shape[0], -1, 1, x.shape[2]*x.shape[3])

            x2 = self.unfold_j(self.pad(x2)).view(x.shape[0], -1, 1, x1.shape[-1])
            # w = self.conv_w(torch.cat([x1, x2], 1)).view(x.shape[0], -1, pow(self.kernel_size, 2), x1.shape[-1])
            # print("x1",x1.size())
            # print("x2",x2.size())
            # batchsize = 4
            # 运算x1*x2
            flage = True
            x1_b1 = x1.split(1, dim=1)
            x2_b1 = x2.split(pow(self.kernel_size, 2), dim=1)
            for i in range(0,len(x1_b1)):
                # out = x1_b1[i] + x2_b1[i]
                out = x1_b1[i] * x2_b1[i]
                if flage:
                    outcat = out
                    flage = False
                else:
                    outcat = torch.cat((outcat, out), 1)
            # print("outcat",outcat.size())

            # 运算x2 +m(x1*x2)
            outcat = x2 + self.m * outcat
            # print("self.m",self.m.size())
            # print("outcat2",outcat.size())
            w = self.conv_w(torch.cat([outcat,x1],1)).view(x.shape[0], -1, pow(self.kernel_size, 2), x1.shape[-1])
            # w = self.conv_w(torch.cat([x1, x2], 1)).view(x.shape[0], -1, pow(self.kernel_size, 2), x1.shape[-1])
            # aggregation: x3 torch.Size([1, 16, 224, 224])
            # aggregation: w  torch.Size([1, 2, 9, 50176])
        # print("aggregation:x3",x3.size())
        # print("aggregation:w",w.size())
        # aggregation--映射函数lamda, 将w映射到x3维度, 矩阵对应位置相乘相加
        x = self.aggregation(x3, w)
        # print("x",x.size())
        return x


class Bottleneck(nn.Module):
    def __init__(self, sa_type, in_planes, rel_planes, mid_planes, out_planes, share_planes=8, kernel_size=7, stride=1):
        super(Bottleneck, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        # print(sa_type, in_planes, rel_planes, mid_planes, share_planes, kernel_size, stride)
        self.sam = SAM(sa_type, in_planes, rel_planes, mid_planes, share_planes, kernel_size, stride)
        self.bn2 = nn.BatchNorm2d(mid_planes)
        self.conv = nn.Conv2d(mid_planes, out_planes, kernel_size=1)
        self.relu = nn.ReLU(inplace=True)
        self.stride = stride

    def forward(self, x):

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

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

        out = self.conv(out)

        out += identity
        return out


class SAN(nn.Module):
    # net = san(sa_type=1, layers=[3, 3, 4, 6, 3], kernels=[3, 7, 7, 7, 7], num_classes=200).cuda().eval()
    def __init__(self, sa_type, block, layers, kernels, num_classes):
        super(SAN, self).__init__()
        # txt.process
        # Embedding 编码方式 in [random, bert, xlnet, albert, word2vec]
        # 加上空格编码，应有71个词
        # self.embed = nn.Embedding(71, 16, padding_idx=0, max_norm=None,
        #                           norm_type=2, scale_grad_by_freq=False, sparse=False)
        self.embed = nn.Embedding(70, 16, padding_idx=0, max_norm=None,
                                  norm_type=2, scale_grad_by_freq=False, sparse=False)  # 70个词, 每个词16维

        self.conv01 = nn.Conv1d(
            in_channels=16,
            out_channels=224,
            kernel_size=3,
            padding=1
        )
        self.conv02 = nn.Conv1d(
            in_channels=224,
            out_channels=448,
            kernel_size=3,
            padding=1
        )
        self.conv_t0 = nn.Conv2d(1, 3, kernel_size=3, stride=1, padding=1)
        # share
        c = 64
        self.conv_in, self.bn_in = conv1x1(3, c), nn.BatchNorm2d(c)
        self.conv0, self.bn0 = conv1x1(c, c), nn.BatchNorm2d(c)
        self.layer0 = self._make_layer(sa_type, block, c, layers[0], kernels[0])

        c *= 4
        self.conv1, self.bn1 = conv1x1(c // 4, c), nn.BatchNorm2d(c)
        self.layer1 = self._make_layer(sa_type, block, c, layers[1], kernels[1])

        c *= 2
        self.conv2, self.bn2 = conv1x1(c // 2, c), nn.BatchNorm2d(c)
        self.layer2 = self._make_layer(sa_type, block, c, layers[2], kernels[2])

        c *= 2
        self.conv3, self.bn3 = conv1x1(c // 2, c), nn.BatchNorm2d(c)
        self.layer3 = self._make_layer(sa_type, block, c, layers[3], kernels[3])

        c *= 2
        self.conv4, self.bn4 = conv1x1(c // 2, c), nn.BatchNorm2d(c)
        self.layer4 = self._make_layer(sa_type, block, c, layers[4], kernels[4])

        self.relu = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc1 = nn.Linear(c, num_classes)

    def _make_layer(self, sa_type, block, planes, blocks, kernel_size=7, stride=1):
        layers = []
        for _ in range(0, blocks):
            layers.append(block(sa_type, planes, planes // 16, planes // 4, planes, 8, kernel_size, stride))
        return nn.Sequential(*layers)

    def forward_share(self, x):
        x = self.relu(self.bn_in(self.conv_in(x)))
        x = self.relu(self.bn0(self.layer0(self.conv0(self.pool(x)))))

        x = self.relu(self.bn1(self.layer1(self.conv1(self.pool(x)))))
        x = self.relu(self.bn2(self.layer2(self.conv2(self.pool(x)))))
        x = self.relu(self.bn3(self.layer3(self.conv3(self.pool(x)))))
        x = self.relu(self.bn4(self.layer4(self.conv4(self.pool(x)))))
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        return x

    def process_txt(self, x):
        # torch.Size([1, 448])
        # x = self.relu(x)
        x = self.embed(x)
        # torch.Size([1, 448, 16])
        x = x.transpose(2, 1)
        # torch.Size([1, 16, 448])
        x = self.conv01(x)
        # torch.Size([1, 224, 448])
        x = self.relu(x)
        x = self.conv02(x)
        # print("conv02", x.size())
        # torch.Size([1, 448, 448])
        x = x.view(-1, 1, 448, 448)
        # torch.Size([1, 1, 448, 448])
        # print("conv0", x.size())
        x = self.conv_t0(x)

        # torch.Size([1, 3, 448, 448])
        return x

    def forward_txt(self, x):
        x = self.process_txt(x)
        x = self.forward_share(x)
        return x


    def forward(self, x1, x2, x3, x):
        x = self.process_txt(x)  # torch.Size([1, 3, 448, 448])
        x = torch.cat((x1, x2, x3, x), dim=0)
        x = self.forward_share(x)
        return x


def san(sa_type, layers, kernels, num_classes):
    model = SAN(sa_type, Bottleneck, layers, kernels, num_classes)
    return model


def get_model_parameters(model):
    total_parameters = 0
    for layer in list(model.parameters()):
        layer_parameter = 1
        for l in list(layer.size()):
            layer_parameter *= l
        total_parameters += layer_parameter
    return total_parameters

if __name__ == '__main__':
    import torch.functional as F
    net = san(sa_type=1, layers=[2, 1, 2, 4, 1], kernels=[3, 7, 7, 7, 7], num_classes=200).cuda().eval()
    # print(get_model_parameters(net))
    img = torch.randn((1, 3, 448, 448)).cuda()
    net.forward_share(img)
    # sa_type, in_planes, rel_planes, out_planes, share_planes, kernel_size=3, stride=1, dilation=1
    # net = SAM(sa_type=1,in_planes=64,rel_planes=4,out_planes=16,share_planes=8,kernel_size=3,stride=1,dilation=1).cuda().eval()
    # print(get_model_parameters(net))
    # sa = torch.randn((1, 64, 224, 224)).cuda()
    # net.forward(sa)



