// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

// Adapter for Gemm in default domain from version 7 to 6

#pragma once

#include <cinttypes>
#include <cstdint>
#include <memory>
#include <vector>

#include "onnx/version_converter/adapters/adapter.h"

namespace ONNX_NAMESPACE::version_conversion {

class Gemm_7_6 final : public Adapter {
 public:
  explicit Gemm_7_6() : Adapter("Gemm", OpSetID(7), OpSetID(6)) {}

  void adapt_gemm_7_6(const std::shared_ptr<Graph>& /*unused*/, Node* node) const {
    const ArrayRef<Value*>& inputs = node->inputs();
    assertInputsAvailable(inputs, name().c_str(), 3);
    const auto& A_shape = inputs[0]->sizes();
    const auto& B_shape = inputs[1]->sizes();
    // Determine if C is broadcastable
    const auto& C_shape = inputs[2]->sizes();
    ONNX_ASSERTM(A_shape.size() == 2, "Gemm input A must have exactly 2 dimensions")
    ONNX_ASSERTM(B_shape.size() == 2, "Gemm input B must have exactly 2 dimensions")
    // Create (M, N) to input to numpy_unibroadcastable
    // TODO(ONNX): Reconcile fact that shapes aren't determined for 1st 2 inputs
    std::vector<Dimension> MN;
    if (node->hasAttribute(ktransA) && node->i(ktransA) == 1) {
      MN.emplace_back(A_shape[1]);
    } else {
      MN.emplace_back(A_shape[0]);
    }
    if (node->hasAttribute(ktransB) && node->i(ktransB) == 1) {
      MN.emplace_back(B_shape[0]);
    } else {
      MN.emplace_back(B_shape[1]);
    }
    int req_broadcast = check_numpy_unibroadcastable_and_require_broadcast(MN, C_shape);
    ONNX_ASSERTM(
        req_broadcast != -1,
        name(),
        " being converted from ",
        static_cast<int64_t>(initial_version().version()),
        " to ",
        static_cast<int64_t>(target_version().version()),
        " does not have broadcastable inputs.")
    if (req_broadcast == 1) {
      node->i_(kbroadcast, 1);
    }
  }

  Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
    adapt_gemm_7_6(graph, node);
    return node;
  }
};

} // namespace ONNX_NAMESPACE::version_conversion