/********************************************************************************
 * 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::TransposeImpl::forward() {
    const Transpose_Op& op = dynamic_cast<const Transpose_Op&>(mOp);
    op.getOutput(0)->copyTranspose(*(op.getInput(0)), op.outputDimsOrder());
}

///////////////////////////////////////////////////

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

Aidge::Transpose_Op::Transpose_Op(const std::vector<Aidge::DimSize_t> &outputDimsOrder)
    : OperatorTensor(Type, {InputCategory::Data}, 1),
        mAttributes(std::make_shared<Attributes_>(
        attr<TransposeAttr::OutputDimsOrder>(outputDimsOrder)))
{
    mImpl = std::make_shared<TransposeImpl>(*this);
}

Aidge::Transpose_Op::Transpose_Op(const Aidge::Transpose_Op& op)
    : OperatorTensor(op),
    mAttributes(op.mAttributes)
{
    if (!op.backend().empty()) {
        SET_IMPL_MACRO(Transpose_Op, *this, op.backend());
    }
    else {
        mImpl = std::make_shared<TransposeImpl>(*this);
    }
}

std::shared_ptr<Aidge::Operator> Aidge::Transpose_Op::clone() const {
    return std::make_shared<Transpose_Op>(*this);
}

bool Aidge::Transpose_Op::forwardDims(bool /*allowDataDependency*/) {
    if (inputsAssociated()) {
        AIDGE_ASSERT(!getInput(0)->empty(), "Not applicable on scalars.");
        // If permutation vector is not given, reverse the dims of input tensor
        if (outputDimsOrder().empty())
        {
            this->outputDimsOrder().resize(getInput(0)->nbDims());
            std::iota(this->outputDimsOrder().rbegin(), this->outputDimsOrder().rend(), 0);
        }

        AIDGE_ASSERT(outputDimsOrder().size() >= getInput(0)->nbDims(),
            "Permutation vector ({}) must have at least the same rank as input tensor ({}).", outputDimsOrder(), getInput(0)->dims());
        std::vector<DimSize_t> outputDims;
        std::size_t i = 0;
        for (; i < getInput(0)->nbDims(); ++i) {
            outputDims.push_back(getInput(0)->dims()[outputDimsOrder()[i]]);
        }
        for (; i < outputDimsOrder().size(); ++i) {
            AIDGE_ASSERT(i == outputDimsOrder()[i],
                "Permutation vector ({}) must be the identity above the input tensor rank ({}).", outputDimsOrder(), getInput(0)->dims());
        }
        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<TransposeImpl>(*this);
    }
    mOutputs[0]->setBackend(name, device);
}

std::set<std::string> Aidge::Transpose_Op::getAvailableBackends() const {
    return Registrar<Transpose_Op>::getKeys();
}

//////////////////////////////////////////////////

std::shared_ptr<Aidge::Node> Aidge::Transpose(const std::vector<Aidge::DimSize_t> &outputDimsOrder,
                                              const std::string& name) {
    return std::make_shared<Node>(std::make_shared<Transpose_Op>(outputDimsOrder), name);
}