import torch
import torch.nn as nn
from torch.nn import init
import functools
from torch.optim import lr_scheduler
import torch.nn.functional as F
from .external_function import SpectralNorm, GroupNorm
import numpy as np


######################################################################################
# base function for network structure
######################################################################################


def init_weights(net, init_type='normal', gain=0.02):
    """Get different initial method for the network weights"""

    def init_func(m):
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:
            init.normal_(m.weight.data, 1.0, 0.02)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.apply(init_func)


def get_norm_layer(norm_type='batch'):
    """Get the normalization layer for the networks"""
    if norm_type == 'batch':
        norm_layer = functools.partial(nn.BatchNorm2d, momentum=0.1, affine=True)
    elif norm_type == 'instance':
        norm_layer = functools.partial(nn.InstanceNorm2d, affine=True)
    elif norm_type == 'group':
        norm_layer = functools.partial(GroupNorm)
    elif norm_type == 'none':
        norm_layer = None
    else:
        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
    return norm_layer


def get_nonlinearity_layer(activation_type='PReLU'):
    """Get the activation layer for the networks"""
    if activation_type == 'ReLU':
        nonlinearity_layer = nn.ReLU()
    elif activation_type == 'SELU':
        nonlinearity_layer = nn.SELU()
    elif activation_type == 'LeakyReLU':
        nonlinearity_layer = nn.LeakyReLU(0.1)
    elif activation_type == 'PReLU':
        nonlinearity_layer = nn.PReLU()
    else:
        raise NotImplementedError('activation layer [%s] is not found' % activation_type)
    return nonlinearity_layer


def get_scheduler(optimizer, opt):
    """Get the training learning rate for different epoch"""
    if opt.lr_policy == 'lambda':
        def lambda_rule(epoch):
            lr_l = 1.0 - max(0, epoch + 1 + 1 + opt.iter_count - opt.niter) / float(opt.niter_decay + 1)
            return lr_l

        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
    elif opt.lr_policy == 'step':
        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
    elif opt.lr_policy == 'exponent':
        scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
    else:
        raise NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
    return scheduler


def print_network(net):
    """print the network"""
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    print(net)
    print('total number of parameters: %.3f M' % (num_params / 1e6))


def init_net(net, init_type='normal', activation='relu', gpu_ids=[]):
    """print the network structure and initial the network"""
    print_network(net)

    if len(gpu_ids) > 0:
        assert (torch.cuda.is_available())
        net.cuda()
        net = torch.nn.DataParallel(net, gpu_ids)
    init_weights(net, init_type)
    return net


def _freeze(*args):
    """freeze the network for forward process"""
    for module in args:
        if module:
            for p in module.parameters():
                p.requires_grad = False


def _unfreeze(*args):
    """ unfreeze the network for parameter update"""
    for module in args:
        if module:
            for p in module.parameters():
                p.requires_grad = True


def spectral_norm(module, use_spect=True):
    """use spectral normal layer to stable the training process"""
    if use_spect:
        return SpectralNorm(module)
    else:
        return module


def coord_conv(input_nc, output_nc, use_spect=False, use_coord=False, with_r=False, use_gated=False, **kwargs):
    """use coord convolution layer to add position information"""
    if use_gated:
        return Gated_Conv(input_nc, output_nc, **kwargs)
    elif use_coord:
        return CoordConv(input_nc, output_nc, with_r, use_spect, **kwargs)
    else:
        return spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)


######################################################################################
# Network basic function
######################################################################################
class AddCoords(nn.Module):
    """
    Add Coords to a tensor
    """

    def __init__(self, with_r=False):
        super(AddCoords, self).__init__()
        self.with_r = with_r

    def forward(self, x):
        """
        :param x: shape (batch, channel, x_dim, y_dim)
        :return: shape (batch, channel+2, x_dim, y_dim)
        """
        B, _, x_dim, y_dim = x.size()

        # coord calculate
        xx_channel = torch.arange(x_dim).repeat(B, 1, y_dim, 1).type_as(x)
        yy_cahnnel = torch.arange(y_dim).repeat(B, 1, x_dim, 1).permute(0, 1, 3, 2).type_as(x)
        # normalization
        xx_channel = xx_channel.float() / (x_dim - 1)
        yy_cahnnel = yy_cahnnel.float() / (y_dim - 1)
        xx_channel = xx_channel * 2 - 1
        yy_cahnnel = yy_cahnnel * 2 - 1

        ret = torch.cat([x, xx_channel, yy_cahnnel], dim=1)

        if self.with_r:
            rr = torch.sqrt(xx_channel ** 2 + yy_cahnnel ** 2)
            ret = torch.cat([ret, rr], dim=1)

        return ret


class CoordConv(nn.Module):
    """
    CoordConv operation
    """

    def __init__(self, input_nc, output_nc, with_r=False, use_spect=False, **kwargs):
        super(CoordConv, self).__init__()
        self.addcoords = AddCoords(with_r=with_r)
        input_nc = input_nc + 2
        if with_r:
            input_nc = input_nc + 1
        self.conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)

    def forward(self, x):
        ret = self.addcoords(x)
        ret = self.conv(ret)

        return ret


class ResBlock(nn.Module):
    """
    Define an Residual block for different types
    """

    def __init__(self, input_nc, output_nc, hidden_nc=None, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(),
                 sample_type='none', use_spect=False, use_coord=False, use_gated=False):
        super(ResBlock, self).__init__()

        hidden_nc = output_nc if hidden_nc is None else hidden_nc
        self.in_nc = input_nc
        self.hi_nc = hidden_nc
        self.ou_nc = output_nc
        self.sample = True
        if sample_type == 'none':
            self.sample = False
        elif sample_type == 'up':
            output_nc = output_nc * 4
            self.pool = nn.PixelShuffle(upscale_factor=2)
        elif sample_type == 'down':
            self.pool = nn.AvgPool2d(kernel_size=2, stride=2)
        else:
            raise NotImplementedError('sample type [%s] is not found' % sample_type)

        kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
        kwargs_short = {'kernel_size': 1, 'stride': 1, 'padding': 0}

        self.conv1 = coord_conv(input_nc, hidden_nc, use_spect, use_coord, use_gated=use_gated, **kwargs)
        self.conv2 = coord_conv(hidden_nc, output_nc, use_spect, use_coord, use_gated=use_gated, **kwargs)
        self.bypass = coord_conv(input_nc, output_nc, use_spect, use_coord, use_gated=use_gated, **kwargs_short)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, self.conv1, nonlinearity, self.conv2, )
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, self.conv1, norm_layer(hidden_nc),
                                       nonlinearity, self.conv2, )

        self.shortcut = nn.Sequential(self.bypass, )

    def forward(self, x):
        if self.sample:
            out = self.pool(self.model(x)) + self.pool(self.shortcut(x))
        else:
            out = self.model(x) + self.shortcut(x)

        return out


class ResBlockEncoderOptimized(nn.Module):
    """
    Define an Encoder block for the first layer of the discriminator and representation network
    """

    def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False,
                 use_coord=False, use_gated=False):
        super(ResBlockEncoderOptimized, self).__init__()

        kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
        kwargs_short = {'kernel_size': 1, 'stride': 1, 'padding': 0}

        self.conv1 = coord_conv(input_nc, output_nc, use_spect, use_coord, use_gated, **kwargs)
        self.conv2 = coord_conv(output_nc, output_nc, use_spect, use_coord, use_gated, **kwargs)
        self.bypass = coord_conv(input_nc, output_nc, use_spect, use_coord, use_gated, **kwargs_short)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(self.conv1, nonlinearity, self.conv2, nn.AvgPool2d(kernel_size=2, stride=2))
        else:
            self.model = nn.Sequential(self.conv1, norm_layer(output_nc), nonlinearity, self.conv2,
                                       nn.AvgPool2d(kernel_size=2, stride=2))

        self.shortcut = nn.Sequential(nn.AvgPool2d(kernel_size=2, stride=2), self.bypass)

    def forward(self, x):
        out = self.model(x) + self.shortcut(x)

        return out


class ResBlockDecoder(nn.Module):
    """
    Define a decoder block
    """

    def __init__(self, input_nc, output_nc, hidden_nc=None, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(),
                 use_spect=False, use_coord=False):
        super(ResBlockDecoder, self).__init__()

        hidden_nc = output_nc if hidden_nc is None else hidden_nc

        self.conv1 = spectral_norm(nn.Conv2d(input_nc, hidden_nc, kernel_size=3, stride=1, padding=1), use_spect)
        self.conv2 = spectral_norm(
            nn.ConvTranspose2d(hidden_nc, output_nc, kernel_size=3, stride=2, padding=1, output_padding=1), use_spect)
        self.bypass = spectral_norm(
            nn.ConvTranspose2d(input_nc, output_nc, kernel_size=3, stride=2, padding=1, output_padding=1), use_spect)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, self.conv1, nonlinearity, self.conv2, )
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, self.conv1, norm_layer(hidden_nc),
                                       nonlinearity, self.conv2, )

        self.shortcut = nn.Sequential(self.bypass)

    def forward(self, x):
        out = self.model(x) + self.shortcut(x)

        return out


class Output(nn.Module):
    """
    Define the output layer
    """

    def __init__(self, input_nc, output_nc, kernel_size=3, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(),
                 use_spect=False, use_coord=False, use_gated=False):
        super(Output, self).__init__()

        kwargs = {'kernel_size': kernel_size, 'padding': 0, 'bias': True}

        self.conv1 = coord_conv(input_nc, output_nc, use_spect, use_coord, use_gated=use_gated, **kwargs)

        if type(norm_layer) == type(None):
            self.model = nn.Sequential(nonlinearity, nn.ReflectionPad2d(int(kernel_size / 2)), self.conv1, nn.Tanh())
        else:
            self.model = nn.Sequential(norm_layer(input_nc), nonlinearity, nn.ReflectionPad2d(int(kernel_size / 2)),
                                       self.conv1, nn.Tanh())

    def forward(self, x):
        out = self.model(x)

        return out


class Auto_Attn(nn.Module):
    """ Short+Long attention Layer"""

    def __init__(self, input_nc, norm_layer=nn.BatchNorm2d):
        super(Auto_Attn, self).__init__()
        self.input_nc = input_nc

        self.query_conv = nn.Conv2d(input_nc, input_nc // 4, kernel_size=1)
        self.gamma = nn.Parameter(torch.zeros(1))
        self.alpha = nn.Parameter(torch.zeros(1))

        self.softmax = nn.Softmax(dim=-1)

        self.model = ResBlock(int(input_nc * 2), input_nc, input_nc, norm_layer=norm_layer, use_spect=True)

    def forward(self, x, pre=None, mask=None):
        """
        inputs :
            x : input feature maps( B X C X W X H)
        returns :
            out : self attention value + input feature
            attention: B X N X N (N is Width*Height)
        """
        # print(x.shape)
        B, C, W, H = x.size()
        proj_query = self.query_conv(x).view(B, -1, W * H)  # B X (N)X C
        proj_key = proj_query  # B X C x (N)

        energy = torch.bmm(proj_query.permute(0, 2, 1), proj_key)  # transpose check
        attention = self.softmax(energy)  # BX (N) X (N)
        proj_value = x.view(B, -1, W * H)  # B X C X N

        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
        out = out.view(B, C, W, H)

        out = self.gamma * out + x

        if type(pre) != type(None):
            # using long distance attention layer to copy information from valid regions
            context_flow = torch.bmm(pre.view(B, -1, W * H), attention.permute(0, 2, 1)).view(B, -1, W, H)
            context_flow = self.alpha * (1 - mask) * context_flow + (mask) * pre
            out = self.model(torch.cat([out, context_flow], dim=1))

        return out, attention


class SimAM(nn.Module):
    def __init__(self, channels=None, e_lambda=1e-4):
        super(SimAM, self).__init__()

        self.activaton = nn.Sigmoid()
        self.e_lambda = e_lambda

    def __repr__(self):
        s = self.__class__.__name__ + '('
        s += ('lambda=%f)' % self.e_lambda)
        return s

    @staticmethod
    def get_module_name():
        return "simam"

    def forward(self, x):
        b, c, h, w = x.size()

        n = w * h - 1

        x_minus_mu_square = (x - x.mean(dim=[2, 3], keepdim=True)).pow(2)
        y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2, 3], keepdim=True) / n + self.e_lambda)) + 0.5

        return x * self.activaton(y)


class HardSPDNorm(nn.Module):
    def __init__(self, n, k, p_input_nc, F_in_nc):
        super(HardSPDNorm, self).__init__()
        self.n = n
        self.k = k
        self.F_in_nc = F_in_nc
        self.gamma_conv = Gated_Conv(2 * p_input_nc, self.F_in_nc, kernel_size=3, stride=1, padding=1)
        self.beta_conv = Gated_Conv(2 * p_input_nc, self.F_in_nc, kernel_size=3, stride=1, padding=1)
        self.dsample_p = nn.AvgPool2d(kernel_size=2, stride=2)
        self.dsample_m = nn.MaxPool2d(kernel_size=2, stride=2)
        self.innorm = nn.InstanceNorm2d(self.F_in_nc, affine=False)

    def forward(self, F_in, img_p, mask, n_ds):
        # print(self.F_in_nc)
        mask = mask.clone()
        # downsample
        for i in range(n_ds):
            img_p = self.dsample_p(img_p)
            mask = self.dsample_m(mask)
        img_p = torch.cat([img_p, mask], dim=1)
        mask, _, _ = torch.chunk(mask, dim=1, chunks=3)
        # D_h
        kernel = torch.ones(mask.shape[1], mask.shape[1], 3, 3).cuda()
        D_h = mask
        msk = D_h.detach()
        msk = torch.where(msk == 1, True, False)
        for i in range(1, self.n + 1):
            D_h = F.conv2d(D_h, kernel, stride=1, padding=1)
            tmp = D_h.detach()
            tmp = torch.where(tmp > 0, True, False)
            tmp = tmp & ~msk
            msk = msk | tmp
            mask[tmp] = 1 / self.k ** (i)
            D_h = mask

        gamma_hp = self.gamma_conv(img_p)
        beta_hp = self.beta_conv(img_p)
        gamma_hd = gamma_hp * D_h
        beta_hd = beta_hp * D_h

        F_in_norm = self.innorm(F_in)
        # print(gamma_hd.shape)
        # print(F_in.shape)
        F_hard = F_in_norm * (gamma_hd + 1) + beta_hd

        return F_hard


class SoftSPDNorm(nn.Module):
    def __init__(self, p_input_nc, F_in_nc):
        super(SoftSPDNorm, self).__init__()
        self.gamma_conv = Gated_Conv(2 * p_input_nc, F_in_nc, kernel_size=3, stride=1, padding=1)
        self.beta_conv = Gated_Conv(2 * p_input_nc, F_in_nc, kernel_size=3, stride=1, padding=1)
        self.p_conv = Gated_Conv(6, F_in_nc, stride=1, kernel_size=1)
        self.f_conv = Gated_Conv(2 * F_in_nc, 1, kernel_size=1, stride=1)
        self.dsample_p = nn.AvgPool2d(kernel_size=2, stride=2)
        self.dsample_m = nn.MaxPool2d(kernel_size=2, stride=2)
        self.sigmoid = nn.Sigmoid()
        self.innorm = nn.InstanceNorm2d(F_in_nc, affine=False)

    def forward(self, F_in, img_p, mask, n_ds):
        mask = mask.clone()
        # downsample
        for i in range(n_ds):
            img_p = self.dsample_p(img_p)
            mask = self.dsample_m(mask)
        img_p = torch.cat([img_p, mask], dim=1)
        mask, _, _ = torch.chunk(mask, dim=1, chunks=3)
        F_in_norm = self.innorm(F_in)

        F_p = self.p_conv(img_p)
        F_mix = torch.cat([F_p, F_in_norm], dim=1)
        F_conv = self.f_conv(F_mix)
        D_s = self.sigmoid(F_conv * (1 - mask) + mask)

        gamma_sp = self.gamma_conv(img_p)
        beta_sp = self.beta_conv(img_p)

        gamma_sd = gamma_sp * D_s
        beta_sd = beta_sp * D_s

        f_in = F_in_norm * gamma_sd
        F_soft = f_in + beta_sd

        return F_soft, mask


class ResBlockSPDNorm(nn.Module):
    def __init__(self, F_in_nc, p_input_nc, n_ds, n, k):
        super(ResBlockSPDNorm, self).__init__()
        self.HardSPDNorm = HardSPDNorm(n, k, p_input_nc, F_in_nc)
        self.SoftSPDNorm = SoftSPDNorm(p_input_nc, F_in_nc)

        self.relu = nn.ReLU()
        self.Conv1 = Gated_Conv(F_in_nc, F_in_nc, stride=1, padding=1, kernel_size=3)
        self.Conv2 = Gated_Conv(F_in_nc, F_in_nc, stride=1, padding=1, kernel_size=3)
        self.Conv3 = Gated_Conv(F_in_nc, F_in_nc, stride=1, padding=1, kernel_size=3)

        self.n_ds = n_ds

    def forward(self, F_in, img_p, mask):
        # the HardSPDNorm
        img_p_re = F.interpolate(
            img_p, size=F_in.size()[2:], mode="nearest"
        )
        mask_re = F.interpolate(mask, size=F_in.size()[2:], mode="nearest")
        out_h1 = self.HardSPDNorm(F_in, img_p_re, mask_re, self.n_ds)
        out_h1 = self.relu(out_h1)
        out_h1 = self.Conv1(out_h1)
        out_h = self.HardSPDNorm(out_h1, img_p_re, mask_re, self.n_ds)
        out_h = self.relu(out_h)
        out_h = self.Conv2(out_h)
        # the SoftSPDNorm
        out_s, mask = self.SoftSPDNorm(F_in, img_p_re, mask_re, self.n_ds)
        out_s = self.relu(out_s)
        out_s = self.Conv3(out_s)
        # output
        out = out_h + out_s

        return out, mask


class Gated_Conv(nn.Module):
    """
    Gated Convolution with spetral normalization
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, batch_norm=False, activation=torch.nn.LeakyReLU(0.2, inplace=True)):
        super(Gated_Conv, self).__init__()
        self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        self.mask_conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        self.activation = activation
        self.batch_norm = batch_norm
        self.batch_norm2d = torch.nn.BatchNorm2d(out_channels)
        self.sigmoid = torch.nn.Sigmoid()
        self.conv2d = torch.nn.utils.spectral_norm(self.conv2d)
        self.mask_conv2d = torch.nn.utils.spectral_norm(self.mask_conv2d)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)

    def gated(self, mask):
        return self.sigmoid(mask)
        #return torch.clamp(mask, -1, 1)

    def forward(self, input):
        x = self.conv2d(input)
        mask = self.mask_conv2d(input)
        if self.activation is not None:
            x = self.activation(x) * self.gated(mask)
        else:
            x = x * self.gated(mask)
        if self.batch_norm:
            return self.batch_norm2d(x)
        else:
            return x


class SPDNormResnetBlock(nn.Module):
    def __init__(self, fin, fout, mask_number, mask_ks):
        super().__init__()
        nhidden = 128
        fmiddle = min(fin, fout)
        lable_nc = 3
        # create conv layers
        self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
        self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
        self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
        self.learned_shortcut = True
        # apply spectral norm if specified
        self.conv_0 = spectral_norm(self.conv_0)
        self.conv_1 = spectral_norm(self.conv_1)
        if self.learned_shortcut:
            self.conv_s = spectral_norm(self.conv_s)
        # define normalization layers
        self.norm_0 = SPDNorm(fin, norm_type="position")
        self.norm_1 = SPDNorm(fmiddle, norm_type="position")
        self.norm_s = SPDNorm(fin, norm_type="position")
        # define the mask weight
        self.v = nn.Parameter(torch.zeros(1))
        self.activeweight = nn.Sigmoid()
        # define the feature and mask process conv
        self.mask_number = mask_number
        self.mask_ks = mask_ks
        pw_mask = int(np.ceil((self.mask_ks - 1.0) / 2))
        self.mask_conv = MaskGet(1, 1, kernel_size=self.mask_ks, padding=pw_mask)
        self.conv_to_f = nn.Sequential(
            Gated_Conv(lable_nc * 2, nhidden, kernel_size=3, padding=1),
            nn.InstanceNorm2d(nhidden),
            nn.ReLU(),
            Gated_Conv(nhidden, fin, kernel_size=3, padding=1),
        )
        self.attention = nn.Sequential(
            nn.Conv2d(fin * 2, fin, kernel_size=3, padding=1), nn.Sigmoid()
        )

    # the semantic feature prior_f form pretrained encoder
    def forward(self, x, prior_image, mask):
        """

        Args:
            x: input feature
            prior_image: the output of PCConv
            mask: mask


        """
        # prepare the forward
        prior_features = []
        b, c, h, w = x.size()
        prior_image_resize = F.interpolate(
            prior_image, size=x.size()[2:], mode="nearest"
        )
        mask_resize = F.interpolate(mask, size=x.size()[2:], mode="nearest")
        prior_image_resize = torch.cat([prior_image_resize, mask_resize], dim=1)
        mask_resize, _, _ = torch.chunk(mask_resize, dim=1, chunks=3)
        # Mask Original and Res path  res weight/ value attention
        prior_feature = self.conv_to_f(prior_image_resize)
        prior_features.append(prior_feature)
        soft_map = self.attention(torch.cat([prior_feature, x], 1))
        # print(soft_map.shape)
        # print(mask_resize.shape)
        soft_map = (1 - mask_resize) * soft_map + mask_resize
        # Mask weight for next process
        mask_pre = mask_resize
        hard_map = 0.0
        for i in range(self.mask_number):
            mask_out = self.mask_conv(mask_pre)
            mask_generate = (mask_out - mask_pre) * (
                1 / (torch.exp(torch.tensor(i + 1).cuda()))
            )
            mask_pre = mask_out
            hard_map = hard_map + mask_generate
        hard_map_inner = (1 - mask_out) * (1 / (torch.exp(torch.tensor(i + 1).cuda())))
        hard_map = hard_map + mask_resize + hard_map_inner
        soft_out = self.conv_s(self.norm_s(x, prior_image_resize, soft_map))
        hard_out = self.conv_0(self.actvn(self.norm_0(x, prior_image_resize, hard_map)))
        hard_out = self.conv_1(
            self.actvn(self.norm_1(hard_out, prior_image_resize, hard_map))
        )
        # Res add
        out = soft_out + hard_out
        return out, prior_features

    def actvn(self, x):
        return F.leaky_relu(x, 2e-1)


def PositionalNorm2d(x, epsilon=1e-5):
    # x: B*C*W*H normalize in C dim
    mean = x.mean(dim=1, keepdim=True)
    std = x.var(dim=1, keepdim=True).add(epsilon).sqrt()
    output = (x - mean) / std
    return output


class SPDNorm(nn.Module):
    def __init__(self, norm_channel, norm_type="batch"):
        super().__init__()
        label_nc = 3
        param_free_norm_type = norm_type
        ks = 3
        if param_free_norm_type == "instance":
            self.param_free_norm = nn.InstanceNorm2d(norm_channel, affine=False)
        elif param_free_norm_type == "batch":
            self.param_free_norm = nn.BatchNorm2d(norm_channel, affine=False)
        elif param_free_norm_type == "position":
            self.param_free_norm = PositionalNorm2d
        else:
            raise ValueError(
                "%s is not a recognized param-free norm type in SPADE"
                % param_free_norm_type
            )

        # The dimension of the intermediate embedding space. Yes, hardcoded.
        pw = ks // 2
        nhidden = 128
        self.mlp_activate = nn.Sequential(
            Gated_Conv(label_nc * 2, nhidden, kernel_size=ks, padding=pw), nn.ReLU()
        )
        self.mlp_gamma = nn.Conv2d(nhidden, norm_channel, kernel_size=ks, padding=pw)
        self.mlp_beta = nn.Conv2d(nhidden, norm_channel, kernel_size=ks, padding=pw)

    def forward(self, x, prior_f, weight):
        normalized = self.param_free_norm(x)
        # Part 2. produce scaling and bias conditioned on condition feature
        actv = self.mlp_activate(prior_f)
        gamma = self.mlp_gamma(actv) * weight
        beta = self.mlp_beta(actv) * weight
        # apply scale and bias
        out = normalized * (1 + gamma) + beta
        return out


class MaskGet(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        padding=0,
        dilation=1,
        groups=1,
        bias=True,
    ):
        super().__init__()
        self.mask_conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            dilation,
            groups,
            False,
        )

        torch.nn.init.constant_(self.mask_conv.weight, 1.0)

    def forward(self, input):
        # depart from partial conv
        # hole region should sed to 0
        with torch.no_grad():
            output_mask = self.mask_conv(input)
        no_update_holes = output_mask == 0
        new_mask = torch.ones_like(output_mask)
        new_mask = new_mask.masked_fill_(no_update_holes.bool(), 0.0)
        return new_mask


class BasicConv(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
        self.relu = nn.ReLU() if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x


class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)


class ChannelGate(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
        super(ChannelGate, self).__init__()
        self.gate_channels = gate_channels
        self.mlp = nn.Sequential(
            Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
            )
        self.pool_types = pool_types

    def forward(self, x):
        channel_att_sum = None
        for pool_type in self.pool_types:
            if pool_type=='avg':
                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( avg_pool )
            elif pool_type=='max':
                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( max_pool )
            elif pool_type=='lp':
                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( lp_pool )
            elif pool_type=='lse':
                # LSE pool only
                lse_pool = logsumexp_2d(x)
                channel_att_raw = self.mlp( lse_pool )

            if channel_att_sum is None:
                channel_att_sum = channel_att_raw
            else:
                channel_att_sum = channel_att_sum + channel_att_raw

        scale = torch.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
        return x * scale


def logsumexp_2d(tensor):
    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
    return outputs


class ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )


class SpatialGate(nn.Module):
    def __init__(self):
        super(SpatialGate, self).__init__()
        kernel_size = 7
        self.compress = ChannelPool()
        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)
    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.spatial(x_compress)
        scale = torch.sigmoid(x_out) # broadcasting
        return x * scale


class CBAM(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
        super(CBAM, self).__init__()
        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
        self.no_spatial = no_spatial
        if not no_spatial:
            self.SpatialGate = SpatialGate()
    def forward(self, x):
        x_out = self.ChannelGate(x)
        if not self.no_spatial:
            x_out = self.SpatialGate(x_out)
        return x_out



class MyAttention(nn.Module):
    """
    Layer implements my version of the self-attention module
    it is mostly same as self attention, but generalizes to
    (k x k) convolutions instead of (1 x 1)
    args:
        in_channels: number of input channels
        out_channels: number of output channels
        activation: activation function to be applied (default: lrelu(0.2))
        kernel_size: kernel size for convolution (default: (1 x 1))
        transpose_conv: boolean denoting whether to use convolutions or transpose
                        convolutions
        squeeze_factor: squeeze factor for query and keys (default: 8)
        stride: stride for the convolutions (default: 1)
        padding: padding for the applied convolutions (default: 1)
        bias: whether to apply bias or not (default: True)
    """

    def __init__(self, in_channels, out_channels,
                 activation=None, kernel_size=(3, 3), transpose_conv=False,
                 use_spectral_norm=True, use_batch_norm=True,
                 squeeze_factor=8, stride=1, padding=1, bias=True):
        """ constructor for the layer """

        from torch.nn import Conv2d, Parameter, \
            Softmax, ConvTranspose2d, BatchNorm2d

        # base constructor call
        super(MyAttention, self).__init__()
        
        # state of the layer
        self.activation = activation
        self.simam = SimAM()
        self.cbam = CBAM(128)
        self.gamma = Parameter(torch.zeros(1))
        self.beta = Parameter(torch.zeros(1))
        self.alpha = Parameter(torch.zeros(1))

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.squeezed_channels = in_channels // squeeze_factor
        self.use_batch_norm = use_batch_norm

        # Modules required for computations
        if transpose_conv:
            self.query_conv = ConvTranspose2d(  # query convolution
                in_channels=in_channels,
                out_channels=in_channels // squeeze_factor,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            )

            self.key_conv = ConvTranspose2d(
                in_channels=in_channels,
                out_channels=in_channels // squeeze_factor,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            )

            self.value_conv = ConvTranspose2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            )

            self.residual_conv = ConvTranspose2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            ) if not use_spectral_norm else SpectralNorm(
                ConvTranspose2d(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=padding,
                    bias=bias
                )
            )

        else:
            self.query_conv = Conv2d(  # query convolution
                in_channels=in_channels,
                out_channels=in_channels // squeeze_factor,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            )

            self.key_conv = Conv2d(
                in_channels=in_channels,
                out_channels=in_channels // squeeze_factor,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            )

            self.value_conv = Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            )

            self.residual_conv = Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias
            ) if not use_spectral_norm else SpectralNorm(
                Conv2d(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=padding,
                    bias=bias
                )
            )

        # softmax module for applying attention
        self.softmax = Softmax(dim=-1)
        self.batch_norm = BatchNorm2d(out_channels)

    def forward(self, x, f):
        """
        forward computations of the layer
        :param x: input feature maps (B x C x H x W)
        :return:
            out: self attention value + input feature (B x O x H x W)
            attention: attention map (B x C x H x W)
        """

        # extract the batch size of the input tensor
        m_batchsize, _, _, _ = x.size()
        # print(f.shape)

        # create the query projection
        proj_query = self.query_conv(x).view(
            m_batchsize, self.squeezed_channels, -1).permute(0, 2, 1)  # B x (N) x C

        # create the key projection
        proj_key = self.key_conv(x).view(
            m_batchsize, self.squeezed_channels, -1)  # B x C x (N)

        # calculate the attention maps
        energy = torch.bmm(proj_query, proj_key)  # energy
        attention = self.softmax(energy)  # attention (B x (N) x (N))

        # create the value projection
        proj_value = self.value_conv(x).view(
            m_batchsize, self.out_channels, -1)  # B X C X N

        # calculate the output
        out = torch.bmm(proj_value, attention.permute(0, 2, 1))

        # calculate the residual output
        res_out = self.residual_conv(x)

        out = out.view(m_batchsize, self.out_channels,
                       res_out.shape[-2], res_out.shape[-1])

        attention = attention.view(m_batchsize, -1,
                                   res_out.shape[-2], res_out.shape[-1])

        if self.use_batch_norm:
            res_out = self.batch_norm(res_out)

        if self.activation is not None:
            out = self.activation(out)
            res_out = self.activation(res_out)

        # apply the residual connections
        # from torchvision import transforms
        sm = self.simam(f)
        # toPIL = transforms.ToPILImage()  # 这个函数可以将张量转为PIL图片，由小数转为0-255之间的像素值
        # print(cbam.chunk(2)[-1].view(128, 32, 32).chunk(128)[-1])
        # pic = toPIL(cbam.chunk(2)[-1].view(128, 32, 32).chunk(128)[-1])
        # pic.save('cbam.jpg')
        # print(cbam)
        # pic = toPIL(sm.chunk(2)[-1].view(128, 32, 32).chunk(128)[-1])
        # pic.save('sm1.jpg')
        sm[:, :, :5, :] = 0
        sm[:, :, :, :5] = 0
        sm[:, :, :, 27:] = 0
        sm[:, :, 27:, :] = 0
        cbam = self.cbam(f)
        # print(cbam.type())
        # print(sm.type())

        map_mean = cbam + sm / 2.0

        # from torchvision import transforms
        # toPIL = transforms.ToPILImage()  # 这个函数可以将张量转为PIL图片，由小数转为0-255之间的像素值
        # print(sm.chunk(2)[-1].view(128, 32, 32).chunk(128)[-1])
        # pic = toPIL(sm.chunk(2)[-1].view(128, 32, 32).chunk(128)[-1])
        # pic.save('map.jpg')
        # out = ((self.gamma * out) + ((1 - self.gamma) * res_out) * (1 - self.beta)) + self.beta * sm
        # out = self.gamma * out + (1 - self.gamma) * res_out
        # out = self.gamma * out + (1 - self.gamma) * sm
        # out = (self.gamma * out + (1 - self.gamma) * res_out) * sm * self.beta
        o1 = self.gamma * out + (1 - self.gamma) * res_out
        out = o1 * self.beta + (1 - self.beta) * map_mean * self.alpha
        # print(cbam)
        # out = self.gamma * out + (1 - self.gamma) * cbam
        return out, attention