/********************************************************************************
 * 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/Transpose.hpp"

#include <cstddef>    // std::size_t
#include <cstdint>    // std::int64_t
#include <memory>
#include <stdexcept>  // std::runtime_error
#include <string>
#include <vector>

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

void Aidge::Transpose_OpImpl::forward() {
    const Transpose_Op& op = dynamic_cast<const Transpose_Op&>(mOp);
    const auto inputDims = op.getInput(0)->dims();
    const auto outputDims = op.getOutput(0)->dims();

    std::vector<std::size_t> outStrides(outputDims.size(), 1);
    for (size_t i = 0; i < outputDims.size(); ++i) {
        for (size_t j = i+1; j < outputDims.size(); ++j)
        {
            outStrides[i] *= outputDims[j];
        }
    }

    std::vector<size_t> indices(outputDims.size(), 0);
    for (size_t i = 0; i < op.getInput(0)->size(); ++i) {
        size_t idx = 0;
        // Permute indices based on OutputDimsOrder attr
        for (int j = outputDims.size() -1; j >=0; --j) {
            idx += indices[op.getAttr<std::vector<DimSize_t>>(0)[j]] * outStrides[j];
        }
        // Copy the value in output
        op.getOutput(0)->getImpl()->copy(op.getInput(0)->getImpl()->rawPtr(i), 1, idx);

        // Update indices for the next iteration
        for (int j = outputDims.size() - 1; j >= 0; --j) {
            if (indices[j] < inputDims[j] - 1) {
                indices[j]++;
                break;
            } else {
                indices[j] = 0;
            }
        }
    }
}

const std::string Aidge::Transpose_Op::Type = "Transpose";

bool Aidge::Transpose_Op::forwardDims(bool /*allowDataDependency*/) {
    // check input has been associated
    if (!getInput(0)) {
        AIDGE_THROW_OR_ABORT(std::runtime_error, "Input was not connected");
    }

    if (!getInput(0)->empty()) {
        const auto& outDimsOrder = getAttr<std::vector<DimSize_t>>(0);
        std::vector<DimSize_t> outputDims;
        for (std::size_t i = 0; i < outDimsOrder.size(); ++i) {
            outputDims.push_back(getInput(0)->dims()[outDimsOrder[i]]);
        }
        mOutputs[0]->resize(outputDims);
        return true;
    }
    return false;
}

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