/********************************************************************************
 * Copyright (c) 2023 CEA-List
 *
 * This program and the accompanying materials are made available under the
 * terms of the Eclipse Public License 2.0 which is available at
 * http://www.eclipse.org/legal/epl-2.0.
 *
 * SPDX-License-Identifier: EPL-2.0
 *
 ********************************************************************************/

#include <memory>
#include <stdexcept>  // std::runtime_error
#include <string>
#include <vector>

#include "aidge/data/Tensor.hpp"
#include "aidge/operator/GlobalAveragePooling.hpp"
#include "aidge/utils/ErrorHandling.hpp"
#include "aidge/utils/Types.h"

const std::string Aidge::GlobalAveragePooling_Op::Type = "GlobalAveragePooling";

Aidge::GlobalAveragePooling_Op::GlobalAveragePooling_Op(const Aidge::GlobalAveragePooling_Op &op)
    : OperatorTensor(op)
{
    if (op.mImpl) {
        SET_IMPL_MACRO(GlobalAveragePooling_Op, *this, op.backend());
    } else {
        mImpl = nullptr;
    }
}

std::shared_ptr<Aidge::Operator> Aidge::GlobalAveragePooling_Op::clone() const {
    return std::make_shared<GlobalAveragePooling_Op>(*this);
}

bool Aidge::GlobalAveragePooling_Op::forwardDims(bool /*allowDataDependency*/) {
    if (inputsAssociated()) {
        AIDGE_ASSERT(getInput(0)->dims().size() >= 3,
                    "GlobalAveragePooling :  needs at least a 3 dimensions input, "
                    "number of input dim : {}",
                    getInput(0)->dims().size());
        // Global average pooling takes each filter, averages its values and uses
        // it as an output(Much like a fancier flatten). 1st dim is batch 2nd is
        // number of filter
        std::vector<DimSize_t> outputDims(getInput(0)->nbDims(), 1);
        outputDims[0] = getInput(0)->dims()[0];
        outputDims[1] = getInput(0)->dims()[1];
        mOutputs[0]->resize(outputDims);
        return true;
    }

    return false;
}

void Aidge::GlobalAveragePooling_Op::setBackend(const std::string &name, Aidge::DeviceIdx_t device) {
    SET_IMPL_MACRO(GlobalAveragePooling_Op, *this, name);
    mOutputs[0]->setBackend(name, device);
}

std::set<std::string> Aidge::GlobalAveragePooling_Op::getAvailableBackends() const {
    return Registrar<GlobalAveragePooling_Op>::getKeys();
}

////////////////////////////////////////

std::shared_ptr<Aidge::Node> Aidge::GlobalAveragePooling(const std::string &name) {
  return std::make_shared<Node>(std::make_shared<GlobalAveragePooling_Op>(), name);
}