/********************************************************************************
 * 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 <vector>
#include <utility>

#include "aidge/backend/OperatorImpl.hpp"
#include "aidge/operator/Slice.hpp"
#include "aidge/utils/Types.h"
#include "aidge/utils/ErrorHandling.hpp"

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

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

    DimSize_t nbAxes = getInput(1)->dims()[0];
    const int* axes = static_cast<const int*>(getInput(1)->getImpl()->rawPtr());
    const int* starts = static_cast<const int*>(getInput(2)->getImpl()->rawPtr());
    const int* ends = static_cast<const int*>(getInput(3)->getImpl()->rawPtr());
    std::vector<DimSize_t> outDims = getInput(0)->dims();
    for(std::size_t i=0; i<nbAxes;++i)
    {
        std::size_t axis = axes[i]>=0?axes[i]:axes[i]+getInput(0)->nbDims();
        outDims[axis] = ends[i] - starts[i] + 1;
    }
    mOutputs[0]->resize(outDims);
}