/*
 * Created by Marcos Luciano
 * https://www.github.com/marcoslucianops
 */

#include "pooling_layer.h"

nvinfer1::ITensor* poolingLayer(
    int layerIdx,
    std::map<std::string, std::string>& block,
    nvinfer1::ITensor* input,
    nvinfer1::INetworkDefinition* network)
{
    nvinfer1::ITensor* output;

    assert(block.at("type") == "maxpool" || block.at("type") == "avgpool");

    if (block.at("type") == "maxpool")
    {
        assert(block.find("size") != block.end());
        assert(block.find("stride") != block.end());

        int size = std::stoi(block.at("size"));
        int stride = std::stoi(block.at("stride"));

        nvinfer1::IPoolingLayer* maxpool
            = network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::Dims{2, {size, size}});
        assert(maxpool != nullptr);
        std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
        maxpool->setName(maxpoolLayerName.c_str());
        maxpool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
        maxpool->setPaddingNd(nvinfer1::Dims{2, {(size - 1) / 2, (size - 1) / 2}});
        if (size == 2 && stride == 1)
        {
            maxpool->setPrePadding(nvinfer1::Dims{2, {0, 0}});
            maxpool->setPostPadding(nvinfer1::Dims{2, {1, 1}});
        }
        output = maxpool->getOutput(0);
    }
    else if (block.at("type") == "avgpool")
    {
        nvinfer1::Dims inputDims = input->getDimensions();
        nvinfer1::IPoolingLayer* avgpool = network->addPoolingNd(
            *input, nvinfer1::PoolingType::kAVERAGE, nvinfer1::Dims{2, {inputDims.d[1], inputDims.d[2]}});
        assert(avgpool != nullptr);
        std::string avgpoolLayerName = "avgpool_" + std::to_string(layerIdx);
        avgpool->setName(avgpoolLayerName.c_str());
        output = avgpool->getOutput(0);
    }
    else
    {
        std::cerr << "Pooling not supported: " << block.at("type") << std::endl;
        std::abort();
    }

    return output;
}