'''
author ：  Hatcher W
Email： hcwang@semi.ac.cn
use: 密度估计网络结构

'''
import torch.nn as nn
from torchvision import  models
from configs.configs import model_head
import torch
import os
import matplotlib.pyplot as plt
import numpy as np
os.environ['CUDA_VISIBLE_DEVICE'] = '5'

class MCNN(nn.Module):
    def __init__(self, load_vgg=True):
        super(MCNN, self).__init__()
        #block1
        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=9, padding=4),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(16),
            nn.MaxPool2d((2, 2))
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=7, padding=3),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(32),
            nn.MaxPool2d((2, 2))
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=16, kernel_size=7, padding=3),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(16),
        )
        self.layer4 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=8, kernel_size=7, padding=3),
            # nn.ReLU(inplace=True),
            # nn.BatchNorm2d(8),
        )
        #block2
        self.layer5 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=20, kernel_size=7, padding=3),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(20),
            nn.MaxPool2d((2, 2))
        )
        self.layer6 = nn.Sequential(
            nn.Conv2d(in_channels=20, out_channels=40, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(40),
            nn.MaxPool2d((2, 2))
        )
        self.layer7 = nn.Sequential(
            nn.Conv2d(in_channels=40, out_channels=20, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(20)
        )
        self.layer8 = nn.Sequential(
            nn.Conv2d(in_channels=20, out_channels=10, kernel_size=5, padding=2),
            # nn.ReLU(inplace=True),
            # nn.BatchNorm2d(10)
        )
        #block3
        self.layer9 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=24, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(24),
            nn.MaxPool2d((2, 2))
        )
        self.layer10 = nn.Sequential(
            nn.Conv2d(in_channels=24, out_channels=48, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(48),
            nn.MaxPool2d((2, 2))
        )
        self.layer11 = nn.Sequential(
            nn.Conv2d(in_channels=48, out_channels=24, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(24)
        )
        self.layer12 = nn.Sequential(
            nn.Conv2d(in_channels=24, out_channels=12, kernel_size=3, padding=1),
            # nn.ReLU(inplace=True),
            # nn.BatchNorm2d(12)
        )
        #end
        self.layer13 = nn.Conv2d(in_channels=30, out_channels=1, kernel_size=1, padding=0)


    def forward(self, x):
        input = x.float()
        x = self.layer1(input)
        x = self.layer2(x)
        x = self.layer3(x)
        out1 = self.layer4(x)

        x = self.layer5(input)
        x = self.layer6(x)
        x = self.layer7(x)
        out2 = self.layer8(x)

        x = self.layer9(input)
        x = self.layer10(x)
        x = self.layer11(x)
        out3 = self.layer12(x)

        out = torch.cat([out1,out2,out3], dim=1)
        out = self.layer13(out)

        return out


    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
class PKSNet(nn.Module):
    def __init__(self, load_vgg=True):
        super(PKSNet, self).__init__()
        #frontend_feat
        self.frontend = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True)
        )
        # left_header
        self.left_header = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(24576, 1000),
            nn.Linear(1000, 5)
        )
        # right_header
        self.right_header = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1),
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(24576, 1000),
            nn.Linear(1000, 13)
        )
        self.bacekend_feat = nn.Sequential(
            #nn.Upsample(scale_factor=2, mode='nearest'),
            #nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            #nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
            #nn.Upsample(scale_factor=2, mode='nearest'),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            #nn.ConvTranspose2d(in_channels=256, out_channels=256, kernel_size=(2 , 2), stride=2),
            #nn.Upsample(scale_factor=2, mode='nearest'),
            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)
        )

        if load_vgg:
            mod = models.vgg16(pretrained=True)
            self._initialize_weights()
            keys1 = list(self.frontend.state_dict().keys())
            keys2 = list(mod.state_dict().keys())
            for i in range(len(keys1)):
                self.frontend.state_dict()[keys1[i]].data[:] = mod.state_dict()[keys2[i]].data[:]
            print('load vgg sucessfully!')


    def forward(self, x):
        x = x.float()
        front_output = self.frontend(x)
        # backend_feat
        density = self.bacekend_feat(front_output)
        # left_header
        if model_head:
            posture = self.left_header(front_output)
        # right_header
            spice = self.right_header(front_output)
            temp = (density, posture, spice)
	        #return density, posture, spice
        else:
            temp = density
        return temp

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
class PKSNet_attention(nn.Module):
    def __init__(self, load_vgg=True):
        super(PKSNet_attention, self).__init__()
        #frontend_feat
        self.frontend = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True)
        )
        # left_header
        self.left_header = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
        )
        self.left_header1 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(24576, 1000),
            nn.Linear(1000, 5)
        )
        # right_header
        self.right_header = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1)
        )
        self.right_header1 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, stride=2, kernel_size=3, padding=1),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(24576, 1000),
            nn.Linear(1000, 13)
        )
        self.bacekend_feat = nn.Sequential(
            #nn.Upsample(scale_factor=2, mode='nearest'),
            #nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            #nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=(2 , 2), stride=2),
            #nn.Upsample(scale_factor=2, mode='nearest'),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            #nn.ConvTranspose2d(in_channels=256, out_channels=256, kernel_size=(2 , 2), stride=2),
            #nn.Upsample(scale_factor=2, mode='nearest'),
            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)
        )

        if load_vgg:
            mod = models.vgg16(pretrained=True)
            self._initialize_weights()
            keys1 = list(self.frontend.state_dict().keys())
            keys2 = list(mod.state_dict().keys())
            for i in range(len(keys1)):
                self.frontend.state_dict()[keys1[i]].data[:] = mod.state_dict()[keys2[i]].data[:]
            print('load vgg sucessfully!')


    def forward(self, x):
        x = x.float()
        front_output = self.frontend(x)
        # backend_feat

        # left_header
        if model_head:
            posture_map = self.left_header(front_output)

            posture = self.left_header1(posture_map)
        # right_header
            spice_map = self.right_header(front_output)
            spice = self.right_header1(spice_map)

            front_output_1 = front_output.clone()
            #print(front_output)
            with torch.set_grad_enabled(False):
                for i in range(4):
                    for j in range(4):
                        #a = front_output[:,:,((i*24)):(((i+1)*24)),((j)*32):(((j+1)*32))]
                        front_output_1[:,:,((i*24)):(((i+1)*24)),((j)*32):(((j+1)*32))]=torch.mul(front_output[:,:,((i*24)):(((i+1)*24)),((j)*32):(((j+1)*32))],spice_map)

            density = self.bacekend_feat(front_output_1)
            temp = (density, posture, spice)
            #print(front_output_1)
	        #return density, posture, spice
        else:
            density = self.bacekend_feat(front_output)
            temp = density
        return temp

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
class CSRNet(nn.Module):
    def __init__(self, load_vgg=True):
        super(MCNN, self).__init__()
        #frontend_feat
        self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]
        self.backend_feat = [512, 512, 512, 256, 128, 64]
        self.frontend = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=1, dilation=1),
            nn.ReLU(inplace=True)
        )
        # backend_feat
        self.layer14 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True))
        self.layer15 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True))
        self.layer16 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True))
        self.layer17 = nn.Sequential(
            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True))
        self.layer18 = nn.Sequential(
            nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True))
        self.layer19 = nn.Sequential(
            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=2, dilation=2),
            nn.ReLU(inplace=True))
        # output_layer
        self.layer20 = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1)

        if load_vgg:
            mod = models.vgg16(pretrained=True)
            self._initialize_weights()
            keys1 = list(self.frontend.state_dict().keys())
            keys2 = list(mod.state_dict().keys())
            for i in range(len(keys1)):
                self.frontend.state_dict()[keys1[i]].data[:] = mod.state_dict()[keys2[i]].data[:]
            print('load vgg sucessfully!')


    def forward(self, x):
        x = x.float()
        x = self.frontend(x)
        x = self.layer14(x)
        x = self.layer15(x)
        x = self.layer16(x)
        x = self.layer17(x)
        x = self.layer18(x)
        x = self.layer19(x)
        x = self.layer20(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
import torch.nn as nn
import torch
from torch.nn import functional as F


class ContextualModule(nn.Module):
    def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
        super(ContextualModule, self).__init__()
        self.scales = []
        self.scales = nn.ModuleList([self._make_scale(features, size) for size in sizes])
        self.bottleneck = nn.Conv2d(features * 2, out_features, kernel_size=1)
        self.relu = nn.ReLU()
        self.weight_net = nn.Conv2d(features,features,kernel_size=1)

    def __make_weight(self,feature,scale_feature):
        weight_feature = feature - scale_feature
        return F.sigmoid(self.weight_net(weight_feature))

    def _make_scale(self, features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
        return nn.Sequential(prior, conv)

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        multi_scales = [F.upsample(input=stage(feats), size=(h, w), mode='bilinear') for stage in self.scales]
        weights = [self.__make_weight(feats,scale_feature) for scale_feature in multi_scales]
        overall_features = [(multi_scales[0]*weights[0]+multi_scales[1]*weights[1]+multi_scales[2]*weights[2]+multi_scales[3]*weights[3])/(weights[0]+weights[1]+weights[2]+weights[3])]+ [feats]
        bottle = self.bottleneck(torch.cat(overall_features, 1))
        return self.relu(bottle)

class CANNet(nn.Module):
    def __init__(self, load_weights=False):
        super(CANNet, self).__init__()
        self.seen = 0
        self.context = ContextualModule(512, 512)
        self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]
        self.backend_feat  = [512, 512, 512,256,128,64]
        self.frontend = make_layers(self.frontend_feat)
        self.backend = make_layers(self.backend_feat,in_channels = 512,batch_norm=True, dilation = True)
        self.output_layer = nn.Conv2d(64, 1, kernel_size=1)


    def forward(self,x):
        x = self.frontend(x)
        x = self.context(x)
        x = self.backend(x)
        x = self.output_layer(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

def make_layers(cfg, in_channels = 3,batch_norm=False,dilation = False):
    if dilation:
        d_rate = 2
    else:
        d_rate = 1
    layers = []
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate,dilation = d_rate)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

if __name__ == '__main__':
    device = torch.device('cuda:2' if torch.cuda.is_available() else 'cpu')
    print(device)
    net = PKSNet_attention().to(device)
    input = torch.randn(1,3,768,1024).to(device)
    print(input)
    density, posture, spice = net(input)
    print(density.shape, posture.shape, spice.shape)