//
//  MobilenetV1.cpp
//  MNN
//
//  Created by MNN on 2020/01/08.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "MobilenetV1.hpp"
#include "Initializer.hpp"
using namespace MNN::Express;

namespace MNN {
namespace Train {
namespace Model {
class _ConvBlock : public Module {
public:
    _ConvBlock(std::vector<int> inputOutputChannels, int stride);

    virtual std::vector<Express::VARP> onForward(const std::vector<Express::VARP> &inputs) override;

    std::shared_ptr<Module> conv3x3;
    std::shared_ptr<Module> bn1;
    std::shared_ptr<Module> conv1x1;
    std::shared_ptr<Module> bn2;
};

std::shared_ptr<Module> ConvBlock(std::vector<int> inputOutputChannels, int stride) {
    return std::shared_ptr<Module>(new _ConvBlock(inputOutputChannels, stride));
}

_ConvBlock::_ConvBlock(std::vector<int> inputOutputChannels, int stride) {
    int inputChannels = inputOutputChannels[0], outputChannels = inputOutputChannels[1];

    NN::ConvOption convOption;
    convOption.kernelSize = {3, 3};
    convOption.channel    = {inputChannels, inputChannels};
    convOption.padMode    = Express::SAME;
    convOption.stride     = {stride, stride};
    convOption.depthwise  = true;
    conv3x3.reset(NN::Conv(convOption, false, std::shared_ptr<Initializer>(Initializer::MSRA())));

    bn1.reset(NN::BatchNorm(inputChannels));

    convOption.reset();
    convOption.kernelSize = {1, 1};
    convOption.channel    = {inputChannels, outputChannels};
    convOption.padMode    = Express::SAME;
    convOption.stride     = {1, 1};
    convOption.depthwise  = false;
    conv1x1.reset(NN::Conv(convOption, false, std::shared_ptr<Initializer>(Initializer::MSRA())));

    bn2.reset(NN::BatchNorm(outputChannels));

    registerModel({conv3x3, bn1, conv1x1, bn2});
}

std::vector<Express::VARP> _ConvBlock::onForward(const std::vector<Express::VARP> &inputs) {
    using namespace Express;
    VARP x = inputs[0];

    x = conv3x3->forward(x);
    x = bn1->forward(x);
    x = _Relu6(x);
    x = conv1x1->forward(x);
    x = bn2->forward(x);
    x = _Relu6(x);

    return {x};
}

MobilenetV1::MobilenetV1(int numClasses, float widthMult, int divisor) {
    NN::ConvOption convOption;
    convOption.kernelSize = {3, 3};
    int outputChannels    = makeDivisible(32 * widthMult, divisor);
    convOption.channel    = {3, outputChannels};
    convOption.padMode    = Express::SAME;
    convOption.stride     = {2, 2};
    conv1.reset(NN::Conv(convOption, false, std::shared_ptr<Initializer>(Initializer::MSRA())));

    bn1.reset(NN::BatchNorm(outputChannels));

    std::vector<std::vector<int> > convSettings;
    convSettings.push_back({64, 1});
    convSettings.push_back({128, 2});
    convSettings.push_back({256, 2});
    convSettings.push_back({512, 6});
    convSettings.push_back({1024, 2});

    int inputChannels = outputChannels;
    for (int i = 0; i < convSettings.size(); i++) {
        auto setting   = convSettings[i];
        outputChannels = setting[0];
        int times      = setting[1];
        outputChannels = makeDivisible(outputChannels * widthMult, divisor);

        for (int j = 0; j < times; j++) {
            int stride = 1;
            if (times > 1 && j == 0) {
                stride = 2;
            }

            convBlocks.emplace_back(ConvBlock({inputChannels, outputChannels}, stride));
            inputChannels = outputChannels;
        }
    }

    dropout.reset(NN::Dropout(0.1));
    fc.reset(NN::Linear(1024, numClasses, true, std::shared_ptr<Initializer>(Initializer::MSRA())));

    registerModel({conv1, bn1, dropout, fc});
    registerModel(convBlocks);
}

std::vector<Express::VARP> MobilenetV1::onForward(const std::vector<Express::VARP> &inputs) {
    using namespace Express;
    VARP x = inputs[0];

    x = conv1->forward(x);
    x = bn1->forward(x);
    x = _Relu6(x);

    for (int i = 0; i < convBlocks.size(); i++) {
        x = convBlocks[i]->forward(x);
    }

    // global avg pooling
    x = _AvePool(x, {-1, -1});

    x = _Convert(x, NCHW);
    x = _Reshape(x, {0, -1});

    x = dropout->forward(x);
    x = fc->forward(x);

    x = _Softmax(x, 1);
    return {x};
}

} // namespace Model
} // namespace Train
} // namespace MNN