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

#include "aidge/operator/MatMul.hpp"
#include "aidge/utils/Types.h"
#include "aidge/utils/ErrorHandling.hpp"

const std::string Aidge::MatMul_Op::Type = "MatMul";

void Aidge::MatMul_Op::computeOutputDims() {
    if (!getInput(0) || !getInput(1)) {
        AIDGE_THROW_OR_ABORT(std::runtime_error, "Missing input. Cannot compute output dimensions for MatMul Operator.");
    }
    if (!getInput(0)->empty() && !getInput(1)->empty())
    {
        const auto dims0 = getInput(0)->dims();
        const auto dims1 = getInput(1)->dims();

        if (dims0.size() > 2 && dims1.size() > 2)
        {
            for (std::size_t d0 = dims0.size()-3, d1 = dims1.size()-3; 
                 (d0>0) && (d1>0);
                 --d0, --d1)
            {
                if(dims0[d0] != dims1[d1])
                    AIDGE_THROW_OR_ABORT(std::runtime_error, "Unsupported sizes for MatMul!");
            }
        }

        std::size_t secondToLastIdx2 = dims1.size()>1 ? dims1.size() - 2 : dims1.size() - 1;
        if(dims0[dims0.size() - 1]  != dims1[secondToLastIdx2])
            AIDGE_THROW_OR_ABORT(std::runtime_error, "Inner dimension missmatch for MatMul!");

        std::vector<std::size_t> outDims;
        if(dims0.size() > 2 || dims1.size() > 2)
        {
            if(dims0.size() > dims1.size())
                std::copy_n(dims0.begin(), dims0.size()-2, std::back_inserter(outDims));
            else
                std::copy_n(dims1.begin(), dims1.size()-2, std::back_inserter(outDims));
        }

        if(dims0.size() > 1)
            outDims.push_back(dims0[dims0.size()-2]);
        if(dims1.size() > 1)
            outDims.push_back(dims1[dims1.size() - 1]);

        mOutputs[0]->resize(outDims);
    }
}