import mxnet as mx
import mxnet.ndarray as nd
import mxnet.gluon as gluon
import mxnet.gluon.nn as nn
import mxnet.autograd as ag

def ConvBlock(channels, kernel_size, strides):
    out = nn.HybridSequential()
    out.add(
        nn.Conv2D(channels, kernel_size, strides=strides, padding=1, use_bias=False),
        nn.BatchNorm(),
        nn.Activation('relu')
    )
    return out

def Conv1x1(channels):
    out = nn.HybridSequential()
    out.add(
        nn.Conv2D(channels, 1, padding=0, use_bias=False),
        nn.BatchNorm(),
        nn.Activation('relu')
    )
    return out

def DWise(channels, stride):
    out = nn.HybridSequential()
    out.add(
        nn.Conv2D(channels, 3, strides=stride, padding=1, groups=channels, use_bias=False),
        nn.BatchNorm(),
        nn.Activation('relu')
    )
    return out

class InvertedResidual(nn.HybridBlock):
    def __init__(self, t, c, s, same_shape=True, **kwargs):
        super(InvertedResidual, self).__init__(**kwargs)
        self.same_shape = same_shape
        self.stride = s
        with self.name_scope(): 
            self.bottleneck = nn.HybridSequential()
            self.bottleneck.add(
                Conv1x1(c*t),
                DWise(c*t, self.stride),
                Conv1x1(c)
            )
            if self.stride == 1 and not self.same_shape:
                self.conv_res = Conv1x1(c)
    def hybrid_forward(self, F, x):
        out = self.bottleneck(x)
        if self.stride == 1:
            if not self.same_shape:
                x = self.conv_res(x)
            out = F.elemwise_add(out, x)
        return out

class MobilenetV2(nn.HybridBlock):
    def __init__(self, num_classes=1000, width_mult=1.0, **kwargs):
        super(MobilenetV2, self).__init__(**kwargs)
        
        self.w = width_mult

        self.b0 = ConvBlock(int(32*self.w), 3, 2)

        self.b1 = nn.HybridSequential()
        self.b1.add(
            InvertedResidual(1, int(16*self.w), 1, same_shape=False)
        )

        self.b2 = nn.HybridSequential()
        self.b2.add(
            InvertedResidual(6, int(24*self.w), 2),
            InvertedResidual(6, int(24*self.w), 2)
        )

        self.b3 = nn.HybridSequential()
        self.b3.add(
            InvertedResidual(6, int(32*self.w), 2),
            InvertedResidual(6, int(32*self.w), 2),
            InvertedResidual(6, int(32*self.w), 2),
        )

        self.b4 = nn.HybridSequential()
        self.b4.add(
            InvertedResidual(6, int(64*self.w), 1, same_shape=False),
            InvertedResidual(6, int(64*self.w), 1),
            InvertedResidual(6, int(64*self.w), 1),
            InvertedResidual(6, int(64*self.w), 1)
        )

        self.b5 = nn.HybridSequential()
        self.b5.add(
            InvertedResidual(6, int(96*self.w), 1, same_shape=False),
            InvertedResidual(6, int(96*self.w), 1),
            InvertedResidual(6, int(96*self.w), 1)
        )

        self.b6 = nn.HybridSequential()
        self.b6.add(
            InvertedResidual(6, int(160*self.w), 2),
            InvertedResidual(6, int(160*self.w), 2),
            InvertedResidual(6, int(160*self.w), 2)
        )

        self.b7 = nn.HybridSequential()
        self.b7.add(
            InvertedResidual(6, int(320*self.w), 1, same_shape=False)
        )

        with self.name_scope():
            self.features = nn.HybridSequential()
            with self.features.name_scope():
                self.features.add(self.b0, self.b1, self.b2, self.b3, self.b4, self.b5, self.b6, self.b7)
                self.features.add(Conv1x1(int(1280*self.w)))
                self.features.add(nn.GlobalAvgPool2D())
                self.features.add(nn.Flatten())
            self.output = nn.Dense(num_classes)
    def hybrid_forward(self, F, x):
        x = self.features(x)
        x = self.output(x)
        return x

net = MobilenetV2(1000,1)

# save as symbol
data =mx.sym.var('data')
sym = net(data)

# plot network graph
mx.viz.plot_network(sym,shape={'data':(8,3,224,224)}, node_attrs={'shape':'oval','fixedsize':'fasl==false'}).view()

