/********************************************************************************
 * 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 "aidge/operator/Slice.hpp"

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>
#include <vector>

#include <fmt/format.h>

#include "aidge/backend/OperatorImpl.hpp"
#include "aidge/data/Data.hpp"
#include "aidge/data/Tensor.hpp"
#include "aidge/utils/ErrorHandling.hpp"
#include "aidge/utils/Types.h"

void Aidge::Slice_OpImpl::forward() {
    const Slice_Op& op = dynamic_cast<const Slice_Op&>(mOp);

    if (!op.getInput(0)) {
        AIDGE_THROW_OR_ABORT(std::runtime_error, "{}: input #0 should be associated with a Tensor", op.Type);
    }
    AIDGE_ASSERT((op.template getAttr<SliceAttr::Starts>().size() == op.template getAttr<SliceAttr::Ends>().size()) &&
                 (op.template getAttr<SliceAttr::Starts>().size() == op.template getAttr<SliceAttr::Axes>().size()),
                 "start, end and axes arguments should be the same size.");

    const auto nbDims = op.getInput(0)->nbDims();
    const auto& inputDims = op.getInput(0)->dims();
    const auto& outputDims = op.getOutput(0)->dims();

    // compute index of the output's first element
    std::size_t beginning = 0;
    const std::size_t nbAxes = op.template getAttr<SliceAttr::Axes>().size();
    for (std::size_t i = 0; i < nbAxes; ++i) {
        // For each slice operation get the params and cast them to size_t
        const DimIdx_t axis = op.template getAttr<SliceAttr::Axes>()[i] >= 0 ?
                            static_cast<DimIdx_t>(op.template getAttr<SliceAttr::Axes>()[i]) :
                            static_cast<DimIdx_t>(op.template getAttr<SliceAttr::Axes>()[i] + static_cast<DimIdx_t>(inputDims.size()));
        const DimSize_t start = op.template getAttr<SliceAttr::Starts>()[i] >= 0 ?
                            static_cast<DimSize_t>(op.template getAttr<SliceAttr::Starts>()[i]) :
                            static_cast<DimSize_t>(op.template getAttr<SliceAttr::Starts>()[i] + static_cast<DimSize_t>(inputDims[axis]));
        const std::size_t stridePostAxis = std::accumulate(inputDims.cbegin()+axis+1, inputDims.cend(), std::size_t(1), std::multiplies<std::size_t>());
        beginning += start * stridePostAxis;
    }

    // for inputDims = {4,5,5,3} & outputDims = {3,2,2,1}: substractDims = {1,5,5,3}
    std::vector<std::size_t> substractedDims = std::vector<std::size_t>(nbDims);
    for (std::size_t i = 0; i < nbDims; ++i) {
        substractedDims[i] = inputDims[i] - outputDims[i];
    }

    // for outputDims = {3,2,2,1}: prodOutputDims = {12,4,2,1}
    std::vector<std::size_t> prodOutputDims = std::vector<std::size_t>(nbDims);
    std::vector<std::size_t> prodInputDims = std::vector<std::size_t>(nbDims + 1);
    prodOutputDims[nbDims - 1] = outputDims[nbDims - 1];
    prodInputDims[nbDims - 1] = inputDims[nbDims - 1];
    prodInputDims[nbDims] = 1;
    for (std::size_t i = 2; i <= nbDims; ++i) {
        prodOutputDims[nbDims - i] = prodOutputDims[nbDims - i + 1] * outputDims[nbDims - i];
        prodInputDims[nbDims - i] = prodInputDims[nbDims - i + 1] * inputDims[nbDims - i];
    }

    std::size_t i = beginning;
    std::size_t size = 0; // number of elements to copy
    std::size_t offset = 0;
    for (std::size_t j = 0; j < prodOutputDims[0];) {
        ++size;
        ++i;
        ++j;
        bool newChunk = false;
        for (std::size_t idx = nbDims - 1; idx > 0; --idx) {
            if (j % prodOutputDims[idx] == 0) {
                i += substractedDims[idx] * prodInputDims[idx + 1];
                newChunk = true;
            }
        }

        if (newChunk) {
            op.getOutput(0)->getImpl()->copy(op.getInput(0)->getImpl()->rawPtr(beginning), size, offset);
            beginning = i;
            offset += size;
            size = 0;
        }
    }

    if (size > 0) {
        op.getOutput(0)->getImpl()->copy(op.getInput(0)->getImpl()->rawPtr(beginning), size, offset);
    }
}

const std::string Aidge::Slice_Op::Type = "Slice";

bool Aidge::Slice_Op::dimsForwarded() const {
    if ((getInput(1) && !getInput(1)->empty())
        || (getInput(2) && !getInput(2)->empty())
        || (getInput(3) && !getInput(3)->empty()))
    {
        // output dims are data dependent
        return false;
    }

    return OperatorTensor::dimsForwarded();
}

bool Aidge::Slice_Op::forwardDims(bool allowDataDependency) {
    // check inputs have been associated
    if (!getInput(0)) {
        AIDGE_THROW_OR_ABORT(std::runtime_error, "{}: input #0 should be associated with a Tensor", type());
    }

    if (getInput(0)->empty()) {
        return false;
    }

    std::shared_ptr<Tensor> fallback;

    if (getInput(1) && !getInput(1)->empty()) {
        if (!this->template getAttr<SliceAttr::Starts>().empty()) {
            Log::notice("Slice_Op: ignoring non-empty Starts attribute because input#1 takes precedence");
        }

        if (!allowDataDependency) {
            Log::warn("Slice_Op: unable to forwardDims() because output dims are data dependent on input#1");
            return false;
        }

        this->template getAttr<SliceAttr::Starts>().clear(); // If both are provided input would override attrs
        this->template getAttr<SliceAttr::Starts>().reserve(getInput(1)->size());
        const auto& starts = getInput(1)->refCastFrom(fallback, NativeType<int64_t>::type, "cpu");
        std::copy_n(static_cast<int64_t*>(starts.getImpl()->hostPtr()),
                    starts.size(),
                    std::back_inserter(this->template getAttr<SliceAttr::Starts>()));
    }

    AIDGE_ASSERT(!this->template getAttr<SliceAttr::Starts>().empty(), "Missing input#1 or Starts attribute");

    if (getInput(2) && !getInput(2)->empty()) {
        if (!this->template getAttr<SliceAttr::Ends>().empty()) {
            Log::notice("Slice_Op: ignoring non-empty Ends attribute because input#2 takes precedence");
        }

        if (!allowDataDependency) {
            Log::warn("Slice_Op: unable to forwardDims() because output dims are data dependent on input#2");
            return false;
        }

        this->template getAttr<SliceAttr::Ends>().clear(); // If both are provided input would override attrs
        this->template getAttr<SliceAttr::Ends>().reserve(getInput(2)->size());
        const auto& ends = getInput(2)->refCastFrom(fallback, NativeType<int64_t>::type, "cpu");
        std::copy_n(static_cast<int64_t*>(ends.getImpl()->hostPtr()),
                    ends.size(),
                    std::back_inserter(this->template getAttr<SliceAttr::Ends>()));
    }

    AIDGE_ASSERT(!this->template getAttr<SliceAttr::Ends>().empty(), "Missing input#2 or Ends attribute");

    if (getInput(3) && !getInput(3)->empty()) {
        if (!this->template getAttr<SliceAttr::Axes>().empty()) {
            Log::notice("Slice_Op: ignoring non-empty Axes attribute because input#3 takes precedence");
        }

        if (!allowDataDependency) {
            Log::warn("Slice_Op: unable to forwardDims() because output dims are data dependent on input#3");
            return false;
        }

        this->template getAttr<SliceAttr::Axes>().clear(); // If both are provided input would override attrs
        this->template getAttr<SliceAttr::Axes>().reserve(getInput(3)->size());
        const auto& axes = getInput(3)->refCastFrom(fallback, NativeType<int8_t>::type, "cpu");
        std::copy_n(static_cast<int8_t*>(axes.getImpl()->hostPtr()),
                    axes.size(),
                    std::back_inserter(this->template getAttr<SliceAttr::Axes>()));
    }

    AIDGE_ASSERT(!this->template getAttr<SliceAttr::Axes>().empty(), "Missing input#3 or Axes attribute");

    const DimSize_t nbAxes = this->template getAttr<SliceAttr::Axes>().size();
    std::vector<DimSize_t> outDims = getInput(0)->dims();
    for (std::size_t i = 0; i < nbAxes; ++i) {
        const DimIdx_t axis = this->template getAttr<SliceAttr::Axes>()[i] >= 0 ?
                        static_cast<DimIdx_t>(this->template getAttr<SliceAttr::Axes>()[i]) :
                        static_cast<DimIdx_t>(this->template getAttr<SliceAttr::Axes>()[i] + static_cast<DimIdx_t>(getInput(0)->nbDims()));
        const DimSize_t start = this->template getAttr<SliceAttr::Starts>()[i] >= 0 ?
                            static_cast<DimSize_t>(this->template getAttr<SliceAttr::Starts>()[i]) :
                            static_cast<DimSize_t>(this->template getAttr<SliceAttr::Starts>()[i] + static_cast<DimSize_t>(getInput(0)->dims()[axis]));
        const DimSize_t end = this->template getAttr<SliceAttr::Ends>()[i] >= 0 ?
                        static_cast<DimSize_t>(this->template getAttr<SliceAttr::Ends>()[i]) :
                        static_cast<DimSize_t>(this->template getAttr<SliceAttr::Ends>()[i] + static_cast<DimSize_t>(getInput(0)->dims()[axis]));

        const std::size_t sliceLength = end - start;
        // Check if slice length is valid
        if (sliceLength > getInput(0)->dims()[axis])
        {
            AIDGE_THROW_OR_ABORT(std::runtime_error, "ROI of Slice operator out of bounds");
        }
        outDims[axis] = sliceLength;
    }
    mOutputs[0]->resize(outDims);
    return true;
}

void Aidge::Slice_Op::setBackend(const std::string& name, Aidge::DeviceIdx_t device) {
    if (Registrar<Slice_Op>::exists({name})){
        SET_IMPL_MACRO(Slice_Op, *this, name);
    }
    else {
        mImpl = std::make_shared<Slice_OpImpl>(*this);
    }
    mOutputs[0]->setBackend(name, device);
}