import torch.nn as nn
from .basic import *

class squeeze_excitation_2d(nn.Module):
    """Squeeze-and-Excitation Block 2D

    Args:
        channel (int): number of input channels.
        channel_reduction (int): channel squeezing factor.
        spatial_reduction (int): pooling factor for x,y axes.
    """
    def __init__(self, channel, channel_reduction=4, spatial_reduction=4, norm_mode='bn', act_mode='elu'):
        super(squeeze_excitation_2d, self).__init__()
        self.pool_size = (spatial_reduction, spatial_reduction)
        layers = [nn.AvgPool2d(kernel_size=self.pool_size, stride=self.pool_size)]
        layers += conv2d_norm_act(channel, channel // channel_reduction, kernel_size=1, padding=0, norm_mode=norm_mode, act_mode=act_mode, return_list=True)
        layers += conv2d_norm_act(channel // channel_reduction, channel, kernel_size=1, padding=0, norm_mode=norm_mode, return_list=True)
        layers = [nn.Sigmoid(),
                nn.Upsample(scale_factor=self.pool_size, mode='trilinear', align_corners=False)]    
        self.se = nn.Sequential(*layers)

    def forward(self, x):
        y = self.se(x)
        z = x + y*x
        return z

class squeeze_excitation_3d(nn.Module):
    """Squeeze-and-Excitation Block 3D

    Args:
        channel (int): number of input channels.
        channel_reduction (int): channel squeezing factor.
        spatial_reduction (int): pooling factor for x,y axes.
        z_reduction (int): pooling factor for z axis.
    """
    def __init__(self, channel, channel_reduction=4, spatial_reduction=4, z_reduction=1, norm_mode='bn', act_mode='elu'):
        super(squeeze_excitation_3d, self).__init__()
        self.pool_size = (z_reduction, spatial_reduction, spatial_reduction)

        layers = [nn.AvgPool3d(kernel_size=self.pool_size, stride=self.pool_size)]
        layers += conv3d_norm_act(channel,  channel//channel_reduction, kernel_size=1, padding=0, norm_mode=norm_mode, act_mode=act_mode, return_list=True)
        layers += conv3d_norm_act(channel//channel_reduction, channel,  kernel_size=1, padding=0, norm_mode=norm_mode, return_list=True)
        layers += [nn.Sigmoid(),
                nn.Upsample(scale_factor=self.pool_size, mode='trilinear', align_corners=False)]
        self.se = nn.Sequential(*layers) 

    def forward(self, x):
        y = self.se(x)
        z = x + y*x
        return z

[((20, 18, 20, 40), 'torch.nn.Sequential', 'nn.Sequential', ({(20, 32, 20, 39): '*layers'}, {}), '(*layers)', True, 'import torch.nn as nn\n'), ((45, 18, 45, 40), 'torch.nn.Sequential', 'nn.Sequential', ({(45, 32, 45, 39): '*layers'}, {}), '(*layers)', True, 'import torch.nn as nn\n'), ((15, 18, 15, 81), 'torch.nn.AvgPool2d', 'nn.AvgPool2d', (), '', True, 'import torch.nn as nn\n'), ((18, 18, 18, 30), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((19, 16, 19, 95), 'torch.nn.Upsample', 'nn.Upsample', (), '', True, 'import torch.nn as nn\n'), ((40, 18, 40, 81), 'torch.nn.AvgPool3d', 'nn.AvgPool3d', (), '', True, 'import torch.nn as nn\n'), ((43, 19, 43, 31), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((44, 16, 44, 95), 'torch.nn.Upsample', 'nn.Upsample', (), '', True, 'import torch.nn as nn\n')]