##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

## Created by: Hang Zhang

## Email: zhanghang0704@gmail.com

## Copyright (c) 2020

##

## LICENSE file in the root directory of this source tree 

##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

import mxnet as mx

from mxnet.gluon import nn

from mxnet.gluon.nn import Conv1D, Block, HybridBlock, Dense, BatchNorm, Activation

__all__ = ['SplitAttentionConv']

USE_BN = True

class SplitAttentionConv(HybridBlock):

    def __init__(self, channels, kernel_size, strides=1, padding=0,

                 dilation=1, groups=1, radix=2, *args, in_channels=None, r=2,

                 norm_layer=BatchNorm, norm_kwargs=None, drop_ratio=0, **kwargs):

        super().__init__()

        norm_kwargs = norm_kwargs if norm_kwargs is not None else {}

        inter_channels = max(in_channels*radix//2//r, 32)

        self.radix = radix

        self.cardinality = groups

        self.conv = Conv1D(channels*radix, kernel_size, strides, padding, dilation,

                           groups=groups*radix, *args, in_channels=in_channels, **kwargs)

        if USE_BN:

            self.bn = norm_layer(in_channels=channels*radix, **norm_kwargs)

        self.relu = Activation('relu')

        self.fc1 = Conv1D(inter_channels, 1, in_channels=channels, groups=self.cardinality)

        if USE_BN:

            self.bn1 = norm_layer(in_channels=inter_channels, **norm_kwargs)

        self.relu1 = Activation('relu')

        if drop_ratio > 0:

            self.drop = nn.Dropout(drop_ratio)

        else:

            self.drop = None

        self.fc2 = Conv1D(channels*radix, 1, in_channels=inter_channels, groups=self.cardinality)

        self.channels = channels

        self.rsoftmax = rSoftMax(radix, groups)

    def hybrid_forward(self, F, x):

        x = self.conv(x)

        if USE_BN:

            x = self.bn(x)

        x = self.relu(x)

        if self.radix > 1:

            splited = F.split(x, self.radix, axis=1)

            gap = sum(splited)

        else:

            gap = x

        gap = F.contrib.AdaptiveAvgPooling1D(gap, 1)

        gap = self.fc1(gap)

        if USE_BN:

            gap = self.bn1(gap)

        atten = self.relu1(gap)

        if self.drop:

            atten = self.drop(atten)

        atten = self.fc2(atten).reshape((0, self.radix, self.channels))

        #atten = self.rsoftmax(atten).reshape((0, -1, 1, 1))

        atten = self.rsoftmax(atten).reshape((0, -1, 1))

        if self.radix > 1:

            atten = F.split(atten, self.radix, axis=1)

            outs = [F.broadcast_mul(att, split) for (att, split) in zip(atten, splited)]

            out = sum(outs)

        else:

            out = F.broadcast_mul(atten, x)

        return out

class rSoftMax(nn.HybridBlock):

    def __init__(self, radix, cardinality):

        super().__init__()

        self.radix = radix

        self.cardinality = cardinality

    def hybrid_forward(self, F, x):

        if self.radix > 1:

            x = x.reshape((0, self.cardinality, self.radix, -1)).swapaxes(1, 2)

            x = F.softmax(x, axis=1)

            x = x.reshape((0, -1))

        else:

            x = F.sigmoid(x)

        return x