# ***************************************************************
# Copyright (c) 2021 Jittor. All Rights Reserved. 
# Maintainers: 
#     Wenyang Zhou <576825820@qq.com>
#     Dun Liang <randonlang@gmail.com>. 
# 
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************
# This model is generated by pytorch converter.

import jittor as jt
from jittor import nn
__all__ = ['MNASNet', 'mnasnet0_5', 'mnasnet0_75', 'mnasnet1_0', 'mnasnet1_3']
_BN_MOMENTUM = (1 - 0.9997)

class _InvertedResidual(nn.Module):

    def __init__(self, in_ch, out_ch, kernel_size, stride, expansion_factor, bn_momentum=0.1):
        super(_InvertedResidual, self).__init__()
        assert (stride in [1, 2])
        assert (kernel_size in [3, 5])
        mid_ch = (in_ch * expansion_factor)
        self.apply_residual = ((in_ch == out_ch) and (stride == 1))
        self.layers = nn.Sequential(nn.Conv(in_ch, mid_ch, 1, bias=False), nn.BatchNorm(mid_ch, momentum=bn_momentum), nn.Relu(), nn.Conv(mid_ch, mid_ch, kernel_size, padding=(kernel_size // 2), stride=stride, groups=mid_ch, bias=False), nn.BatchNorm(mid_ch, momentum=bn_momentum), nn.Relu(), nn.Conv(mid_ch, out_ch, 1, bias=False), nn.BatchNorm(out_ch, momentum=bn_momentum))

    def execute(self, input):
        if self.apply_residual:
            return (self.layers(input) + input)
        else:
            return self.layers(input)

def _stack(in_ch, out_ch, kernel_size, stride, exp_factor, repeats, bn_momentum):
    assert (repeats >= 1)
    first = _InvertedResidual(in_ch, out_ch, kernel_size, stride, exp_factor, bn_momentum=bn_momentum)
    remaining = []
    for _ in range(1, repeats):
        remaining.append(_InvertedResidual(out_ch, out_ch, kernel_size, 1, exp_factor, bn_momentum=bn_momentum))
    return nn.Sequential(first, *remaining)

def _round_to_multiple_of(val, divisor, round_up_bias=0.9):
    assert (0.0 < round_up_bias < 1.0)
    new_val = max(divisor, ((int((val + (divisor / 2))) // divisor) * divisor))
    return (new_val if (new_val >= (round_up_bias * val)) else (new_val + divisor))

def _get_depths(alpha):
    depths = [24, 40, 80, 96, 192, 320]
    return [_round_to_multiple_of((depth * alpha), 8) for depth in depths]

class MNASNet(nn.Module):
    """ MNASNet model architecture. version=2.

    Args:
    
    * alpha: Depth multiplier.
    * num_classes: Number of classes. Default: 1000.
    * dropout: Dropout probability of dropout layer.
    """
    _version = 2

    def __init__(self, alpha, num_classes=1000, dropout=0.2):
        super(MNASNet, self).__init__()
        assert (alpha > 0.0)
        self.alpha = alpha
        self.num_classes = num_classes
        depths = _get_depths(alpha)
        layers = [
            nn.Conv(3, 32, 3, padding=1, stride=2, bias=False), 
            nn.BatchNorm(32, momentum=_BN_MOMENTUM), 
            nn.Relu(), 
            nn.Conv(32, 32, 3, padding=1, stride=1, groups=32, bias=False), 
            nn.BatchNorm(32, momentum=_BN_MOMENTUM), 
            nn.Relu(), 
            nn.Conv(32, 16, 1, padding=0, stride=1, bias=False), 
            nn.BatchNorm(16, momentum=_BN_MOMENTUM), 
            _stack(16, depths[0], 3, 2, 3, 3, _BN_MOMENTUM),
            _stack(depths[0], depths[1], 5, 2, 3, 3, _BN_MOMENTUM),
            _stack(depths[1], depths[2], 5, 2, 6, 3, _BN_MOMENTUM),
            _stack(depths[2], depths[3], 3, 1, 6, 2, _BN_MOMENTUM),
            _stack(depths[3], depths[4], 5, 2, 6, 4, _BN_MOMENTUM),
            _stack(depths[4], depths[5], 3, 1, 6, 1, _BN_MOMENTUM),
            nn.Conv(depths[5], 1280, 1, padding=0, stride=1, bias=False), 
            nn.BatchNorm(1280, momentum=_BN_MOMENTUM), 
            nn.Relu()
        ]
        self.layers = nn.Sequential(*layers)
        self.classifier = nn.Sequential(nn.Dropout(p=dropout), nn.Linear(1280, num_classes))

    def execute(self, x):
        x = self.layers(x)
        x = x.mean([2, 3])
        return self.classifier(x)

def mnasnet0_5(pretrained=False, **kwargs):
    model = MNASNet(0.5, **kwargs)
    if pretrained: model.load("jittorhub://mnasnet0_5.pkl")
    return model

def mnasnet0_75(pretrained=False, **kwargs):
    model = MNASNet(0.75, **kwargs)
    if pretrained: model.load("jittorhub://mnasnet0_75.pkl")
    return model

def mnasnet1_0(pretrained=False, **kwargs):
    model = MNASNet(1.0, **kwargs)
    if pretrained: model.load("jittorhub://mnasnet1_0.pkl")
    return model

def mnasnet1_3(pretrained=False, **kwargs):
    model = MNASNet(1.3, **kwargs)
    if pretrained: model.load("jittorhub://mnasnet1_3.pkl")
    return model
