/********************************************************************************
 * 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 <numeric>    // std::accumulate
#include <functional> // std::multiplies

#include "aidge/operator/Slice.hpp"

#include "aidge/backend/cpu/operator/SliceImpl.hpp"
#include "aidge/backend/cpu/operator/SliceImpl_forward_kernels.hpp"
#include "aidge/utils/Types.h"
#include <vector>
#include <cassert>
#include <tuple>

Aidge::NbElts_t Aidge::SliceImpl_cpu::getNbRequiredData(const Aidge::IOIndex_t /*inputIdx*/) const {
    assert(std::static_pointer_cast<Tensor>(mOp.getRawInput(0)) && "requires valid input");

    // Requires the whole tensors
    const auto& inputDims = std::static_pointer_cast<Tensor>(mOp.getRawInput(0))->dims();

    return std::accumulate(inputDims.begin(), inputDims.end(), static_cast<NbElts_t>(1),
                            std::multiplies<NbElts_t>());
}

Aidge::NbElts_t Aidge::SliceImpl_cpu::getNbRequiredProtected(const Aidge::IOIndex_t /*inputIdx*/) const { return 0; }

Aidge::NbElts_t Aidge::SliceImpl_cpu::getRequiredMemory(const Aidge::IOIndex_t outputIdx,
                            const std::vector<Aidge::DimSize_t>& inputsSize) const {
    (void)outputIdx;
    (void)inputsSize;
    const auto& outputDims = std::static_pointer_cast<Tensor>(mOp.getRawOutput(0))->dims();
    return std::accumulate(outputDims.begin(), outputDims.end(), static_cast<NbElts_t>(1),
                            std::multiplies<NbElts_t>());
}

Aidge::NbElts_t Aidge::SliceImpl_cpu::getNbConsumedData(const Aidge::IOIndex_t /*inputIdx*/) const {
    return mNbConsumedData[0];
}

Aidge::NbElts_t Aidge::SliceImpl_cpu::getNbProducedData(const Aidge::IOIndex_t /*outputIdx*/) const {
    return mNbProducedData[0];
}

void Aidge::SliceImpl_cpu::updateConsummerProducer() {
    // each input is consumed by the minimum amount for a forward pass
    mNbConsumedData[0] += getNbRequiredData(0);

    mNbProducedData[0] += getRequiredMemory(0, {});
}

void Aidge::SliceImpl_cpu::forward() {
    // FIXME: uncomment the following code once memory handling will work
    assert(std::static_pointer_cast<Tensor>(mOp.getRawInput(0)) && "missing input #0");

    // Find the correct kernel type
    auto kernelFunc = Registrar<SliceImplForward_cpu>::create(
            {std::static_pointer_cast<Tensor>(mOp.getRawInput(0))->dataType()});

    // Call kernel
    kernelFunc(dynamic_cast<const Slice_Op&>(mOp).getStaticAttributes(),
            std::static_pointer_cast<Tensor>(mOp.getRawInput(0))->dims(),
            std::static_pointer_cast<Tensor>(mOp.getRawInput(0))->getImpl()->rawPtr(),
            std::static_pointer_cast<Tensor>(mOp.getRawOutput(0))->getImpl()->rawPtr()
            );

    // each input is consumed by the minimum amount for a forward pass
    mNbConsumedData[0] += getNbRequiredData(0);

    mNbProducedData[0] += getRequiredMemory(0, {});
}

void Aidge::SliceImpl_cpu::backward() { fmt::print("Not implemented yet.\n"); }