#include <algorithm>

#include "IndexingUtils.hpp"

using namespace mlir;

namespace onnx_mlir {

int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }

FailureOr<int64_t> normalizeAxisChecked(int64_t axis, int64_t rank) {
  int64_t normalizedAxis = normalizeAxis(axis, rank);
  if (normalizedAxis < 0 || normalizedAxis >= rank)
    return failure();
  return normalizedAxis;
}

int64_t normalizeIndex(int64_t index, int64_t dimSize) { return index >= 0 ? index : dimSize + index; }

static SmallVector<int64_t> normalizeAxesImpl(std::optional<ArrayAttr> axesAttr, int64_t rank) {
  SmallVector<int64_t> normalizedAxes;
  if (!axesAttr) {
    normalizedAxes.reserve(rank);
    for (int64_t axis = 0; axis < rank; ++axis)
      normalizedAxes.push_back(axis);
  }
  else {
    normalizedAxes.reserve(axesAttr->size());
    for (Attribute attr : *axesAttr)
      normalizedAxes.push_back(normalizeAxis(cast<IntegerAttr>(attr).getInt(), rank));
    llvm::sort(normalizedAxes);
    normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
  }
  return normalizedAxes;
}

FailureOr<SmallVector<int64_t>> normalizeAxesChecked(std::optional<ArrayAttr> axesAttr, int64_t rank) {
  SmallVector<int64_t> normalizedAxes = normalizeAxesImpl(axesAttr, rank);
  for (int64_t axis : normalizedAxes)
    if (axis < 0 || axis >= rank)
      return failure();
  return normalizedAxes;
}

} // namespace onnx_mlir