import torch.nn as nn
import torch.nn.functional as F
import torch

try:
    from .sync_batchnorm import SynchronizedBatchNorm3d
except:
    pass


def normalization(planes, norm='bn'):
    if norm == 'bn':
        m = nn.BatchNorm3d(planes)
    elif norm == 'gn':
        m = nn.GroupNorm(4, planes)
    elif norm == 'in':
        m = nn.InstanceNorm3d(planes)
    elif norm == 'sync_bn':
        m = SynchronizedBatchNorm3d(planes)
    else:
        raise ValueError('normalization type {} is not supported'.format(norm))
    return m


class Conv3d_Block(nn.Module):
    def __init__(self, num_in, num_out, kernel_size=1, stride=1, g=1, padding=None, norm=None):
        super(Conv3d_Block, self).__init__()
        if padding == None:
            padding = (kernel_size - 1) // 2
        self.bn = normalization(num_in, norm=norm)
        self.act_fn = nn.ReLU(inplace=True)
        self.conv = nn.Conv3d(num_in, num_out, kernel_size=kernel_size, padding=padding, stride=stride, groups=g,
                              bias=False)

    def forward(self, x):  # BN + Relu + Conv
        h = self.act_fn(self.bn(x))
        h = self.conv(h)
        return h


class DilatedConv3DBlock(nn.Module):
    def __init__(self, num_in, num_out, kernel_size=(1, 1, 1), stride=1, g=1, d=(1, 1, 1), norm=None):
        super(DilatedConv3DBlock, self).__init__()
        assert isinstance(kernel_size, tuple) and isinstance(d, tuple)

        padding = tuple(
            [(ks - 1) // 2 * dd for ks, dd in zip(kernel_size, d)]
        )

        self.bn = normalization(num_in, norm=norm)
        self.act_fn = nn.ReLU(inplace=True)
        self.conv = nn.Conv3d(num_in, num_out, kernel_size=kernel_size, padding=padding, stride=stride, groups=g,
                              dilation=d, bias=False)

    def forward(self, x):
        h = self.act_fn(self.bn(x))
        h = self.conv(h)
        return h


class MFunit(nn.Module):
    def __init__(self, num_in, num_out, g=1, stride=1, d=(1, 1), norm=None):
        """  The second 3x3x1 group conv is replaced by 3x3x3.
        :param num_in: number of input channels
        :param num_out: number of output channels
        :param g: groups of group conv.
        :param stride: 1 or 2
        :param d: tuple, d[0] for the first 3x3x3 conv while d[1] for the 3x3x1 conv
        :param norm: Batch Normalization
        """
        super(MFunit, self).__init__()
        num_mid = num_in if num_in <= num_out else num_out
        self.conv1x1x1_in1 = Conv3d_Block(num_in, num_in // 4, kernel_size=1, stride=1, norm=norm)
        self.conv1x1x1_in2 = Conv3d_Block(num_in // 4, num_mid, kernel_size=1, stride=1, norm=norm)
        self.conv3x3x3_m1 = DilatedConv3DBlock(num_mid, num_out, kernel_size=(3, 3, 3), stride=stride, g=g,
                                               d=(d[0], d[0], d[0]), norm=norm)  # dilated
        self.conv3x3x3_m2 = DilatedConv3DBlock(num_out, num_out, kernel_size=(3, 3, 1), stride=1, g=g,
                                               d=(d[1], d[1], 1), norm=norm)
        # self.conv3x3x3_m2 = DilatedConv3DBlock(num_out,num_out,kernel_size=(1,3,3),stride=1,g=g,d=(1,d[1],d[1]),norm=norm)

        # skip connection
        if num_in != num_out or stride != 1:
            if stride == 1:
                self.conv1x1x1_shortcut = Conv3d_Block(num_in, num_out, kernel_size=1, stride=1, padding=0, norm=norm)
            if stride == 2:
                # if MF block with stride=2, 2x2x2
                self.conv2x2x2_shortcut = Conv3d_Block(num_in, num_out, kernel_size=2, stride=2, padding=0,
                                                       norm=norm)  # params

    def forward(self, x):
        x1 = self.conv1x1x1_in1(x)
        x2 = self.conv1x1x1_in2(x1)
        x3 = self.conv3x3x3_m1(x2)
        x4 = self.conv3x3x3_m2(x3)

        shortcut = x

        if hasattr(self, 'conv1x1x1_shortcut'):
            shortcut = self.conv1x1x1_shortcut(shortcut)
        if hasattr(self, 'conv2x2x2_shortcut'):
            shortcut = self.conv2x2x2_shortcut(shortcut)

        return x4 + shortcut


class DMFUnit(nn.Module):
    # weighred add
    def __init__(self, num_in, num_out, g=1, stride=1, norm=None, dilation=None):
        super(DMFUnit, self).__init__()
        self.weight1 = nn.Parameter(torch.ones(1))
        self.weight2 = nn.Parameter(torch.ones(1))
        self.weight3 = nn.Parameter(torch.ones(1))

        num_mid = num_in if num_in <= num_out else num_out

        self.conv1x1x1_in1 = Conv3d_Block(num_in, num_in // 4, kernel_size=1, stride=1, norm=norm)
        self.conv1x1x1_in2 = Conv3d_Block(num_in // 4, num_mid, kernel_size=1, stride=1, norm=norm)

        self.conv3x3x3_m1 = nn.ModuleList()
        if dilation == None:
            dilation = [1, 2, 3]
        for i in range(3):
            self.conv3x3x3_m1.append(
                DilatedConv3DBlock(num_mid, num_out, kernel_size=(3, 3, 3), stride=stride, g=g,
                                   d=(dilation[i], dilation[i], dilation[i]), norm=norm)
            )

        # It has not Dilated operation
        self.conv3x3x3_m2 = DilatedConv3DBlock(num_out, num_out, kernel_size=(3, 3, 1), stride=(1, 1, 1), g=g,
                                               d=(1, 1, 1), norm=norm)
        # self.conv3x3x3_m2 = DilatedConv3DBlock(num_out, num_out, kernel_size=(1, 3, 3), stride=(1,1,1), g=g,d=(1,1,1), norm=norm)

        # skip connection
        if num_in != num_out or stride != 1:
            if stride == 1:
                self.conv1x1x1_shortcut = Conv3d_Block(num_in, num_out, kernel_size=1, stride=1, padding=0, norm=norm)
            if stride == 2:
                self.conv2x2x2_shortcut = Conv3d_Block(num_in, num_out, kernel_size=2, stride=2, padding=0, norm=norm)

    def forward(self, x):
        x1 = self.conv1x1x1_in1(x)
        x2 = self.conv1x1x1_in2(x1)
        x3 = self.weight1 * self.conv3x3x3_m1[0](x2) + self.weight2 * self.conv3x3x3_m1[1](x2) + self.weight3 * \
             self.conv3x3x3_m1[2](x2)
        x4 = self.conv3x3x3_m2(x3)
        shortcut = x
        if hasattr(self, 'conv1x1x1_shortcut'):
            shortcut = self.conv1x1x1_shortcut(shortcut)
        if hasattr(self, 'conv2x2x2_shortcut'):
            shortcut = self.conv2x2x2_shortcut(shortcut)
        return x4 + shortcut


class MFNet(nn.Module):  #
    # [96]   Flops:  13.361G  &  Params: 1.81M
    # [112]  Flops:  16.759G  &  Params: 2.46M
    # [128]  Flops:  20.611G  &  Params: 3.19M
    def __init__(self, c=4, n=32, channels=128, groups=16, norm='bn', num_classes=4, att=False, **kwargs):
        super(MFNet, self).__init__()

        self.att = att
        # Entry flow
        self.encoder_block1 = nn.Conv3d(c, n, kernel_size=3, padding=1, stride=2, bias=False)  # H//2
        self.encoder_block2 = nn.Sequential(
            MFunit(n, channels, g=groups, stride=2, norm=norm),  # H//4 down
            MFunit(channels, channels, g=groups, stride=1, norm=norm),
            MFunit(channels, channels, g=groups, stride=1, norm=norm)
        )
        #
        self.encoder_block3 = nn.Sequential(
            MFunit(channels, channels * 2, g=groups, stride=2, norm=norm),  # H//8
            MFunit(channels * 2, channels * 2, g=groups, stride=1, norm=norm),
            MFunit(channels * 2, channels * 2, g=groups, stride=1, norm=norm)
        )

        self.encoder_block4 = nn.Sequential(  # H//8,channels*4
            MFunit(channels * 2, channels * 3, g=groups, stride=2, norm=norm),  # H//16
            MFunit(channels * 3, channels * 3, g=groups, stride=1, norm=norm),
            MFunit(channels * 3, channels * 2, g=groups, stride=1, norm=norm),
        )

        self.upsample1 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=False)  # H//8
        self.decoder_block1 = MFunit(channels * 2 + channels * 2, channels * 2, g=groups, stride=1, norm=norm)

        self.upsample2 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=False)  # H//4
        self.decoder_block2 = MFunit(channels * 2 + channels, channels, g=groups, stride=1, norm=norm)

        self.upsample3 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=False)  # H//2
        self.decoder_block3 = MFunit(channels + n, n, g=groups, stride=1, norm=norm)
        self.upsample4 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=False)  # H
        self.seg = nn.Conv3d(n, num_classes, kernel_size=1, padding=0, stride=1, bias=False)

        self.softmax = nn.Softmax(dim=1)
        if self.att:
            self.cor_att0 = CornerAttention(channels, **kwargs)#to verify
            self.cor_att1 = CornerAttention(channels, **kwargs)
            self.cor_att2 = CornerAttention(channels, **kwargs)

        # Initialization
        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                torch.nn.init.torch.nn.init.kaiming_normal_(m.weight)  #
            elif isinstance(m, nn.BatchNorm3d) or isinstance(m, nn.GroupNorm) or isinstance(m, SynchronizedBatchNorm3d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        # Encoder
        x1 = self.encoder_block1(x)  # H//2 down
        x2 = self.encoder_block2(x1)  # H//4 down
        x3 = self.encoder_block3(x2)  # H//8 down
        x4 = self.encoder_block4(x3)  # H//16
        # Decoder
        y1 = self.upsample1(x4)  # H//8

        y1 = torch.cat([x3, y1], dim=1)
        y1 = self.decoder_block1(y1)

        y2 = self.upsample2(y1)  # H//4

        y2 = torch.cat([x2, y2], dim=1)
        y2 = self.decoder_block2(y2)

        y3 = self.upsample3(y2)  # H//2

        y3 = torch.cat([x1, y3], dim=1)
        y3 = self.decoder_block3(y3)

        y4 = self.upsample4(y3)
        y4 = self.seg(y4)
        if hasattr(self, 'softmax'):
            y4 = self.softmax(y4)
        return y4


class DMFNet(MFNet):  # softmax
    # [128]  Flops:  27.045G  &  Params: 3.88M
    def __init__(self, c=4, n=32, channels=128, groups=16, norm='bn', num_classes=4):
        super(DMFNet, self).__init__(c, n, channels, groups, norm, num_classes)

        self.encoder_block2 = nn.Sequential(
            DMFUnit(n, channels, g=groups, stride=2, norm=norm, dilation=[1, 2, 3]),  # H//4 down
            DMFUnit(channels, channels, g=groups, stride=1, norm=norm, dilation=[1, 2, 3]),  # Dilated Conv 3
            DMFUnit(channels, channels, g=groups, stride=1, norm=norm, dilation=[1, 2, 3])
        )

        self.encoder_block3 = nn.Sequential(
            DMFUnit(channels, channels * 2, g=groups, stride=2, norm=norm, dilation=[1, 2, 3]),  # H//8
            DMFUnit(channels * 2, channels * 2, g=groups, stride=1, norm=norm, dilation=[1, 2, 3]),  # Dilated Conv 3
            DMFUnit(channels * 2, channels * 2, g=groups, stride=1, norm=norm, dilation=[1, 2, 3])
        )


class CornerAttention(nn.Module):
    def __init__(self, in_channels, shape=(64, 64, 64), pooling='avg', reduction=4, wchannel=None, middle_norm='sig',
                 **kwargs):
        super(CornerAttention, self).__init__()
        middle_ch = in_channels // reduction
        self.shape = shape
        self.pooling = pooling
        self.wchannel = wchannel
        self.middle_norm = middle_norm

        self.hconv1 = nn.Sequential(
            nn.InstanceNorm3d(in_channels),
            nn.Conv3d(in_channels=in_channels, out_channels=middle_ch, kernel_size=1, stride=1, padding=0),
            nn.ReLU(inplace=True),
            nn.Dropout3d(p=0.5),
            nn.InstanceNorm3d(in_channels),
            nn.Conv3d(in_channels=middle_ch, out_channels=middle_ch, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=middle_ch, out_channels=in_channels, kernel_size=1, stride=1, padding=0),
        )

        self.wconv1 = nn.Sequential(
            nn.InstanceNorm3d(in_channels),
            nn.Conv3d(in_channels=in_channels, out_channels=middle_ch, kernel_size=1, stride=1, padding=0),
            nn.ReLU(inplace=True),
            nn.Dropout3d(p=0.5),
            nn.InstanceNorm3d(in_channels),
            nn.Conv3d(in_channels=middle_ch, out_channels=middle_ch, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=middle_ch, out_channels=in_channels, kernel_size=1, stride=1, padding=0),
        )

        self.dconv1 = nn.Sequential(
            nn.InstanceNorm3d(in_channels),
            nn.Conv3d(in_channels=in_channels, out_channels=middle_ch, kernel_size=1, stride=1, padding=0),
            nn.ReLU(inplace=True),
            nn.Dropout3d(p=0.5),
            nn.InstanceNorm3d(in_channels),
            nn.Conv3d(in_channels=middle_ch, out_channels=middle_ch, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=middle_ch, out_channels=in_channels, kernel_size=1, stride=1, padding=0),
        )
        if self.pooling == 'max':
            self.dpooling = nn.AdaptiveMaxPool3d((shape[0], 1, 1))
            self.hpooling = nn.AdaptiveMaxPool3d((1, shape[1], 1))
            self.wpooling = nn.AdaptiveMaxPool3d((1, 1, shape[1]))
        else:
            self.dpooling = nn.AdaptiveAvgPool3d((shape[0], 1, 1))
            self.hpooling = nn.AdaptiveAvgPool3d((1, shape[1], 1))
            self.wpooling = nn.AdaptiveAvgPool3d((1, 1, shape[1]))

        self.dconv2 = nn.Conv3d(in_channels, in_channels, (3, 1, 1), stride=(1, 1, 1), padding=(1, 0, 0))
        self.hconv2 = nn.Conv3d(in_channels, in_channels, (1, 3, 1), stride=(1, 1, 1), padding=(0, 1, 0))
        self.wconv2 = nn.Conv3d(in_channels, in_channels, (1, 1, 3), stride=(1, 1, 1), padding=(0, 0, 1))

        self.fuse = nn.Conv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
        self.sig = nn.Sigmoid()

    def forward(self, x):

        # pdb.set_trace()
        dx = self.dconv1(x)
        hx = self.hconv1(x)
        wx = self.wconv1(x)

        if self.pooling == 'max':
            dx = F.adaptive_max_pool3d(dx, (x.shape[2], 1, 1))
            hx = F.adaptive_max_pool3d(hx, (1, x.shape[3], 1))
            wx = F.adaptive_max_pool3d(wx, (1, 1, x.shape[4]))
        else:
            dx = F.adaptive_avg_pool3d(dx, (x.shape[2], 1, 1))
            hx = F.adaptive_avg_pool3d(hx, (1, x.shape[3], 1))
            wx = F.adaptive_avg_pool3d(wx, (1, 1, x.shape[4]))

        # dx = self.dpooling(dx)
        # hx = self.hpooling(hx)
        # wx = self.wpooling(wx)

        # dx = self.sig(self.dconv2(dx))
        # hx = self.sig(self.hconv2(hx))
        # wx = self.sig(self.wconv2(wx))

        if self.middle_norm == 'sig':
            dx = self.sig(self.dconv2(dx))
            hx = self.sig(self.hconv2(hx))
            wx = self.sig(self.wconv2(wx))
        else:
            dx = F.softmax(self.dconv2(dx), 1)
            hx = F.softmax(self.hconv2(hx), 1)
            wx = F.softmax(self.wconv2(wx), 1)

        attx = dx * hx * wx

        attx = self.sig(self.fuse(attx))
        x = x * attx
        return x


if __name__ == '__main__':
    import os

    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    x = torch.rand((1, 4, 128, 128, 128))  # [bsize,channels,Height,Width,Depth]
    model = DMFNet(c=4, groups=16, norm='sync_bn', num_classes=4)

    y = model(x)
    print(y.shape)
