/********************************************************************************
 * 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 <cassert>
#include <cstddef>
#include <string>
#include <vector>

#include "aidge/operator/Gather.hpp"
#include "aidge/utils/Types.h"
#include "aidge/utils/ErrorHandling.hpp"

const std::string Aidge::Gather_Op::Type = "Gather";

void Aidge::Gather_Op::computeOutputDims() {
    // check inputs have been associated
    if (!getInput(0) || !getInput(1)) {
        AIDGE_THROW_OR_ABORT(std::runtime_error, "At least one input was not connected");
    }

    if (getInput(1)->nbDims()!=2){
        AIDGE_THROW_OR_ABORT(std::runtime_error, "Indices input must be a 2D Tensor");
    }

    std::vector<DimSize_t> outDims = getInput(0)->dims();
    std::vector<DimSize_t> indexesDims = getInput(1)->dims();
    int axisIdx = this->template getAttr<GatherAttr::Axis>()>=0?this->template getAttr<GatherAttr::Axis>():this->template getAttr<GatherAttr::Axis>()+outDims.size();
    outDims.erase(outDims.begin() + static_cast<std::size_t>(axisIdx));
    outDims.insert(outDims.begin() + static_cast<std::size_t>(axisIdx), indexesDims.begin(),indexesDims.end());
    mOutputs[0]->resize(outDims);
}