#include "mlir/IR/ValueRange.h"

#include "llvm/ADT/STLExtras.h"

#include <cassert>

#include "Common.hpp"
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"

using namespace llvm;
using namespace mlir;

namespace onnx_mlir {

FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anchor, mlir::Value value) {
  auto byteSize = pim::getCheckedShapedTypeSizeInBytes(cast<ShapedType>(value.getType()), anchor, "tensor byte size");
  if (failed(byteSize))
    return failure();
  return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
}

Operation* getEarliestUserWithinBlock(mlir::Value value) {
  auto users = value.getUsers();

  assert(!users.empty());

  Operation* earliestUser = *users.begin();
  for (auto curUser : users)
    if (curUser->isBeforeInBlock(earliestUser))
      earliestUser = curUser;

  return earliestUser;
}

SmallVector<mlir::Value> getOpOperandsSortedByUses(Operation* operation) {
  auto operandsAndUses =
    map_to_vector(operation->getOperands(), [](mlir::Value operand) -> std::pair<mlir::Value, size_t> {
      return {operand, std::distance(operand.use_begin(), operand.use_end())};
    });
  sort(operandsAndUses, [](auto a, auto b) { return a.second < b.second; });
  return map_to_vector(operandsAndUses, [](auto operandAndUse) { return operandAndUse.first; });
}

bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
  for (Operation* user : value.getUsers()) {
    if (user->getBlock() != operation->getBlock())
      return true;
    if (operation->isBeforeInBlock(user))
      return true;
  }
  return false;
}

mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Operation* operation) {
  assert("Only support operations with a single result" && operation->getNumResults() == 1);
  mlir::Value result = operation->getResult(0);
  auto resultType = result.getType();
  assert("Only support result ShapedType as result type" && isa<ShapedType>(resultType));

  SmallVector<mlir::Value> operands = getOpOperandsSortedByUses(operation);
  auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) {
    return operand.getType() == resultType && !hasLaterUserInBlock(operand, operation);
  });
  auto bestOperand = validOperands.begin();

  if (bestOperand != validOperands.end())
    return *bestOperand;

  auto resultShapedType = cast<ShapedType>(resultType);
  rewriter.setInsertionPoint(operation);
  return tensor::EmptyOp::create(
    rewriter, operation->getLoc(), resultShapedType.getShape(), resultShapedType.getElementType());
}

} // namespace onnx_mlir