import torch
from torch import nn
from torch.nn import init
import torch.nn.functional as F
from torch.optim import Adam
import numpy
from einops import rearrange
import time
from transformer import Transformer
from Intra_MLP import index_points,knn_l2

# vgg choice
base = {'vgg': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']}

# vgg16
def vgg(cfg, i=3, batch_norm=True):
    layers = []
    in_channels = i
    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=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return layers


def hsp(in_channel, out_channel):
    layers = nn.Sequential(nn.Conv2d(in_channel, out_channel, 1, 1),
                           nn.ReLU())
    return layers

def cls_modulation_branch(in_channel, hiden_channel):
    layers = nn.Sequential(nn.Linear(in_channel, hiden_channel),
                           nn.ReLU())
    return layers

def cls_branch(hiden_channel, class_num):
    layers = nn.Sequential(nn.Linear(hiden_channel, class_num),
                           nn.Sigmoid())
    return layers

def intra():
    layers = []
    layers += [nn.Conv2d(512, 512, 1, 1)]
    layers += [nn.Sigmoid()]
    return layers

def concat_r():
    layers = []
    layers += [nn.Conv2d(512, 512, 1, 1)]
    layers += [nn.ReLU()]
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    layers += [nn.ReLU()]
    layers += [nn.ConvTranspose2d(512, 512, 4, 2, 1)]
    return layers

def concat_1():
    layers = []
    layers += [nn.Conv2d(512, 512, 1, 1)]
    layers += [nn.ReLU()]
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    layers += [nn.ReLU()]
    return layers

def mask_branch():
    layers = []
    layers += [nn.Conv2d(512, 2, 3, 1, 1)]
    layers += [nn.ConvTranspose2d(2, 2, 8, 4, 2)]
    layers += [nn.Softmax2d()]
    return layers

def incr_channel():
    layers = []
    layers += [nn.Conv2d(128, 512, 3, 1, 1)]
    layers += [nn.Conv2d(256, 512, 3, 1, 1)]
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    return layers

def incr_channel2():
    layers = []
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    layers += [nn.Conv2d(512, 512, 3, 1, 1)]
    layers += [nn.ReLU()]
    return layers

def norm(x, dim):
    squared_norm = (x ** 2).sum(dim=dim, keepdim=True)
    normed = x / torch.sqrt(squared_norm)
    return normed

def fuse_hsp(x, p,group_size=5):
    
    t = torch.zeros(group_size, x.size(1))
    for i in range(x.size(0)):
        tmp = x[i, :]
        if i == 0:
            nx = tmp.expand_as(t)
        else:
            nx = torch.cat(([nx, tmp.expand_as(t)]), dim=0)
    nx = nx.view(x.size(0)*group_size, x.size(1), 1, 1)
    y = nx.expand_as(p)
    return y


class Model(nn.Module):
    def __init__(self, device, base, incr_channel, incr_channel2, hsp1, hsp2, cls_m, cls, concat_r, concat_1, mask_branch, intra,demo_mode=False):
        super(Model, self).__init__()
        self.base = nn.ModuleList(base)
        self.sp1 = hsp1
        self.sp2 = hsp2
        self.cls_m = cls_m
        self.cls = cls
        self.incr_channel1 = nn.ModuleList(incr_channel)
        self.incr_channel2 = nn.ModuleList(incr_channel2)
        self.concat4 = nn.ModuleList(concat_r)
        self.concat3 = nn.ModuleList(concat_r)
        self.concat2 = nn.ModuleList(concat_r)
        self.concat1 = nn.ModuleList(concat_1)
        self.mask = nn.ModuleList(mask_branch)
        self.extract = [13, 23, 33, 43]
        self.device = device
        self.group_size = 5
        self.intra = nn.ModuleList(intra)
        self.transformer_1=Transformer(512,4,4,782,group=self.group_size)
        self.transformer_2=Transformer(512,4,4,782,group=self.group_size)
        self.demo_mode=demo_mode

    def forward(self, x):
        # backbone, p is the pool2, 3, 4, 5
        p = list()
        for k in range(len(self.base)):
            x = self.base[k](x)
            if k in self.extract:
                p.append(x)
        
        
        # increase the channel
        newp = list()
        newp_T=list()
        for k in range(len(p)):
            np = self.incr_channel1[k](p[k])
            np = self.incr_channel2[k](np)
            newp.append(self.incr_channel2[4](np))
            if k==3:
              tmp_newp_T3=self.transformer_1(newp[k])
              newp_T.append(tmp_newp_T3)
            if k==2:
              newp_T.append(self.transformer_2(newp[k]))
            if k<2:
              newp_T.append(None)


        # intra-MLP
        point = newp[3].view(newp[3].size(0), newp[3].size(1), -1)
        point = point.permute(0,2,1)

        idx = knn_l2(self.device, point, 4, 1)
        feat=idx
        new_point = index_points(self.device, point,idx)

        group_point = new_point.permute(0, 3, 2, 1)
        group_point = self.intra[0](group_point)
        group_point = torch.max(group_point, 2)[0]  # [B, D', S]

        intra_mask = group_point.view(group_point.size(0), group_point.size(1), 7, 7)
        intra_mask = intra_mask + newp[3]
        
        spa_mask = self.intra[1](intra_mask)


        x = newp[3]
        x = self.sp1(x)
        x = x.view(-1, x.size(1), x.size(2) * x.size(3))
        x = torch.bmm(x, x.transpose(1, 2))
        x = x.view(-1, x.size(1) * x.size(2))
        x = x.view(x.size(0) // self.group_size, x.size(1), -1, 1)
        x = self.sp2(x)
        x = x.view(-1, x.size(1), x.size(2) * x.size(3))
        x = torch.bmm(x, x.transpose(1, 2))
        x = x.view(-1, x.size(1) * x.size(2))

        #cls pred
        cls_modulated_vector = self.cls_m(x)
        cls_pred = self.cls(cls_modulated_vector)

        #semantic and spatial modulator
        g1 = fuse_hsp(cls_modulated_vector, newp[0],self.group_size)
        g2 = fuse_hsp(cls_modulated_vector, newp[1],self.group_size)
        g3 = fuse_hsp(cls_modulated_vector, newp[2],self.group_size)
        g4 = fuse_hsp(cls_modulated_vector, newp[3],self.group_size)

        spa_1 = F.interpolate(spa_mask, size=[g1.size(2), g1.size(3)], mode='bilinear')
        spa_1 = spa_1.expand_as(g1)
        spa_2 = F.interpolate(spa_mask, size=[g2.size(2), g2.size(3)], mode='bilinear')
        spa_2 = spa_2.expand_as(g2)
        spa_3 = F.interpolate(spa_mask, size=[g3.size(2), g3.size(3)], mode='bilinear')
        spa_3 = spa_3.expand_as(g3)
        spa_4 = F.interpolate(spa_mask, size=[g4.size(2), g4.size(3)], mode='bilinear')
        spa_4 = spa_4.expand_as(g4)
        
        y4 = newp_T[3] * g4 + spa_4
        for k in range(len(self.concat4)):
            y4 = self.concat4[k](y4)
        
        y3 = newp_T[2] * g3 + spa_3
        
        for k in range(len(self.concat3)):
            y3 = self.concat3[k](y3)
            if k == 1:
                y3 = y3 + y4
        
        y2 = newp[1] * g2 + spa_2
        
        #print(y2.shape)
        
        for k in range(len(self.concat2)):
            y2 = self.concat2[k](y2)
            if k == 1:
                y2 = y2 + y3
        y1 = newp[0] * g1 + spa_1

        for k in range(len(self.concat1)):
            y1 = self.concat1[k](y1)
            if k == 1:
                y1 = y1 + y2
        y = y1
        if self.demo_mode:
            tmp=F.interpolate(y1, size=[14,14], mode='bilinear')
            tmp=tmp.permute(0,2,3,1).contiguous().reshape(tmp.shape[0]*tmp.shape[2]*tmp.shape[3],tmp.shape[1])
            tmp=tmp/torch.norm(tmp,p=2,dim=1).unsqueeze(1)
            feat2=(tmp@tmp.t())
            feat=F.interpolate(y, size=[14,14], mode='bilinear')
        
        # decoder
        for k in range(len(self.mask)):
            
            y = self.mask[k](y)
        mask_pred = y[:, 0, :, :]
        if self.demo_mode:
            return cls_pred, mask_pred,feat,feat2
        else:
            return cls_pred, mask_pred



# build the whole network
def build_model(device,demo_mode=False):
    return Model(device,
                 vgg(base['vgg']),
                 incr_channel(),
                 incr_channel2(),
                 hsp(512, 64),
                 hsp(64**2, 32),
                 cls_modulation_branch(32**2, 512),
                 cls_branch(512, 78),
                 concat_r(),
                 concat_1(),
                 mask_branch(),
                 intra(),demo_mode)

# weight init
def xavier(param):
    init.xavier_uniform_(param)

def weights_init(m):
    if isinstance(m, nn.Conv2d):
        xavier(m.weight.data)
    elif isinstance(m, nn.BatchNorm2d):
        init.constant_(m.weight, 1)
        init.constant_(m.bias, 0)

'''import os
os.environ['CUDA_VISIBLE_DEVICES']='6'
gpu_id='cuda:0'
device = torch.device(gpu_id)
nt=build_model(device).to(device)
it=2
bs=1
gs=5
sum=0
with torch.no_grad():
  for i in range(it):
    A=torch.rand(bs*gs,3,448,256).cuda()
    A=A*2-1
    start=time.time()
    nt(A)
    sum+=time.time()-start
print(sum/bs/gs/it)'''