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

// Helper Methods for Adapters

#include "onnx_lib/version_converter/helper.h"

#include <vector>

namespace ONNX_LIGHT_NAMESPACE::version_conversion {
int check_numpy_unibroadcastable_and_require_broadcast(const std::vector<Dimension> &input1_sizes,
                                                       const std::vector<Dimension> &input2_sizes) {
  // Check that input1 is larger
  if (input1_sizes.size() < input2_sizes.size())
    return -1;
  // Check that axis is input1_sizes.size()-input2_sizes.size()
  bool broadcast = false;
  auto axis = input1_sizes.size() - input2_sizes.size();
  for (size_t i = 0; i < input2_sizes.size(); i++) {
    if (input2_sizes[i].dim != input1_sizes[axis + i].dim && input2_sizes[i].dim != 1)
      return -1;
    if (input2_sizes[i].dim != input1_sizes[axis + i].dim)
      broadcast = true;
  }
  // Return true if broadcasting is required
  if (input1_sizes.size() > input2_sizes.size() || broadcast)
    return 1;
  else
    return 0;
}

void assert_numpy_multibroadcastable(const std::vector<Dimension> &input1_sizes,
                                     const std::vector<Dimension> &input2_sizes) {
  // Generalize above for multibroadcastable case
  const std::vector<Dimension> *A_ptr = &input1_sizes;
  const std::vector<Dimension> *B_ptr = &input2_sizes;
  int A = 1;
  int B = 2;
  if (input1_sizes.size() < input2_sizes.size()) {
    A_ptr = &input2_sizes;
    B_ptr = &input1_sizes;
    A = 2;
    B = 1;
  }
  const std::vector<Dimension> &A_sizes = *A_ptr;
  const std::vector<Dimension> &B_sizes = *B_ptr;
  auto axis = A_sizes.size() - B_sizes.size();
  for (size_t i = 0; i < B_sizes.size(); i++) {
    ONNX_ASSERTM(B_sizes[i].dim == A_sizes[axis + i].dim || B_sizes[i].dim == 1 ||
                     A_sizes[axis + i].dim == 1,
                 "Dimension ", i, " of input ", B, " does not match dimension ", axis + i,
                 " of input ", A, ", and neither's value is 1")
  }
}

void assertNotParams(const std::vector<Dimension> &sizes) {
  for (const Dimension &dim : sizes) {
    ONNX_ASSERTM(dim.is_int, dim.param, " Dimension is a param instead of an int.")
  }
}

void assertInputsAvailable(const ArrayRef<Value *> &inputs, const char *name, uint64_t num_inputs) {
  ONNX_ASSERTM(inputs.size() == num_inputs, name, " in opset version 6 can only broadcast between ",
               num_inputs, " inputs")
  for (size_t i = 0; i < num_inputs; i++) {
    ONNX_ASSERTM(inputs[i]->has_sizes(), "Shape of input ", i, " is not available.")
    assertNotParams(inputs[i]->sizes());
  }
}
} // namespace ONNX_LIGHT_NAMESPACE::version_conversion
