#ifndef LIB_UTILS_UTILS_H_
#define LIB_UTILS_UTILS_H_

#include <cstdint>
#include <functional>
#include <optional>
#include <queue>
#include <string>
#include <utility>
#include <vector>

#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/include/mlir/IR/Attributes.h"            // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/IR/Dialect.h"          // from @llvm-project
#include "mlir/include/mlir/IR/Location.h"         // from @llvm-project
#include "mlir/include/mlir/IR/Operation.h"        // from @llvm-project
#include "mlir/include/mlir/IR/Types.h"            // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"            // from @llvm-project
#include "mlir/include/mlir/IR/Visitors.h"         // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"        // from @llvm-project
#include "mlir/include/mlir/Support/WalkResult.h"  // from @llvm-project

namespace mlir {
namespace heir {

// Helper for the "overloaded" pattern
template <class... Ts>
struct Overloaded : Ts... {
  using Ts::operator()...;
};
template <class... Ts>
Overloaded(Ts...) -> Overloaded<Ts...>;

using OpPredicate = std::function<bool(Operation*)>;
using IsValidTypeFn = std::function<LogicalResult(const Type&)>;
using IsValidValueFn = std::function<LogicalResult(const Value&)>;
using ValueProcessor = std::function<void(const Value&)>;

using TypePredicate = std::function<bool(const Type&)>;

using DialectPredicate = std::function<bool(Dialect*)>;

using IndexTupleConsumer = std::function<void(const std::vector<int64_t>&)>;

template <typename... OpTys>
OpPredicate OpEqual() {
  return [](Operation* op) { return mlir::isa<OpTys...>(op); };
}

template <typename... TypeTys>
TypePredicate TypeEqual() {
  return [](const Type& type) { return mlir::isa<TypeTys...>(type); };
}

template <typename... DialectTys>
DialectPredicate DialectEqual() {
  return [](Dialect* dialect) { return mlir::isa<DialectTys...>(dialect); };
}

// Walks the given op, applying the predicate to traversed ops until the
// predicate returns true, then returns the operation that matched, or
// nullptr if there were no matches.
Operation* walkAndDetect(Operation* op, OpPredicate predicate);

// specialization for detecting a specific operation type
template <typename... OpTys>
bool containsAnyOperations(Operation* op) {
  Operation* foundOp = walkAndDetect(op, OpEqual<OpTys...>());
  return foundOp != nullptr;
}

/// Apply isValidType to the operands and results, returning an appropriate
/// logical result.
LogicalResult validateTypes(Operation* op, IsValidTypeFn isValidType);

/// Apply isValidType to the operands, returning an appropriate logical result.
LogicalResult validateValues(Operation* op, IsValidValueFn isValidValue);

/// Walk the IR and apply a predicate to all argument values
/// encountered, returning failure if any type is invalid. Invalidity is
/// determined by whether the IsValidTypeFn returns a failed LogicalResult. If
/// err is provided, report and error to the user when the first invalid value
/// is encountered.
LogicalResult walkAndValidateValues(
    Operation* op, IsValidValueFn isValidValue,
    std::optional<std::string> err = std::nullopt);

namespace detail {
template <typename T>
Operation* getOperation(T& op) {
  return op.getOperation();
}

template <>
inline Operation* getOperation<Operation*>(Operation*& op) {
  return op;
}
}  // namespace detail

/// Walk the IR and apply a predicate to all SSA values and produced values.
/// Values will be visited once for the operation that defines them, and once
/// for each use. The valueProcessor must be aware that it may be called
/// multiple times for the same value.
void walkValues(Operation* op, ValueProcessor valueProcessor);

/// Walk the IR and apply a predicate to all argument and result types
/// encountered, returning failure if any type is invalid. Invalidity is
/// determined by whether the IsValidTypeFn returns a failed LogicalResult.
/// If err is provided, report and error to the user when the first invalid type
/// is encountered.
template <typename OpTy>
LogicalResult walkAndValidateTypes(
    Operation* op, IsValidTypeFn isValidType,
    std::optional<std::string> err = std::nullopt) {
  LogicalResult res = success();
  op->walk([&](OpTy op) {
    Operation* operation = detail::getOperation(op);
    res = validateTypes(operation, isValidType);
    if (failed(res) && err.has_value()) operation->emitError() << err.value();
    return failed(res) ? WalkResult::interrupt() : WalkResult::advance();
  });
  return res;
}

/// Get the function op that a call op is calling.
FailureOr<func::FuncOp> getCalledFunction(func::CallOp callOp);

// Helper to traverse op and any nested functions called.
// Applies callback to matching ops of type OpType in func, and in any funcs
// called by func. The walk is interrupted if callback returns interrupt().
template <typename FnT>
WalkResult walkFuncAndCallees(func::FuncOp func, FnT&& callback) {
  std::queue<func::FuncOp> worklist;
  worklist.push(func);
  DenseSet<func::FuncOp> done = {};
  while (!worklist.empty()) {
    func::FuncOp next = worklist.front();
    worklist.pop();
    WalkResult result = next.walk([&](Operation* op) {
      if (auto callOp = dyn_cast<func::CallOp>(op)) {
        auto maybeCallee = getCalledFunction(callOp);
        if (succeeded(maybeCallee) && !done.contains(maybeCallee.value()) &&
            maybeCallee.value() != next)
          worklist.push(maybeCallee.value());
      }
      return callback(op);
    });
    if (result.wasInterrupted()) return WalkResult::interrupt();
    done.insert(next);
  }
  return WalkResult::advance();
}

// Returns true if the op contains ops from the given dialects.
template <typename... Dialects>
bool containsDialects(Operation* op) {
  Operation* foundOp = walkAndDetect(op, [&](Operation* op) {
    return DialectEqual<Dialects...>()(op->getDialect());
  });
  return foundOp != nullptr;
}

// Returns true if the op contains argument values of the given type.
// NOTE: any_of instead of all_of
bool containsArgumentOfType(Operation* op, TypePredicate predicate);

template <typename... TypeTys>
bool containsArgumentOfType(Operation* op) {
  return containsArgumentOfType(op, TypeEqual<TypeTys...>());
}

// A helper to iterate over the space of indices of a multidimensional tensor
// whose shape is given by `shape`, passing each index tuple visited to
// `process`.
//
// If fixedIndices and fixedIndexValues are nonempty, iterate over the
// remaining indices and always populate the index tuple provided to
// `process` with these fixed index values.
//
// E.g., if shape is {2, 3, 4}, fixedIndices is {1}, and fixedIndexValues is
// {2}, then this will iterate over dimensions 0 and 2 in the usual order, but
// dimension 1 will always be 2.
void iterateIndices(ArrayRef<int64_t> shape, const IndexTupleConsumer& process,
                    ArrayRef<int64_t> fixedIndices = {},
                    ArrayRef<int64_t> fixedIndexValues = {});

/// Convert a double to a string with 2 decimal precision.
std::string doubleToString2Prec(double value);

// Get a scalar or dense attribute given the type and value.
TypedAttr getScalarOrDenseAttr(Type tensorOrScalarType, APInt value);

// Get a scalar or dense attribute given the type and value.
TypedAttr getScalarOrDenseAttr(Type tensorOrScalarType, APFloat value);

Operation* makeAppropriatelyTypedAddOp(OpBuilder& builder, Location loc,
                                       Value lhs, Value rhs,
                                       ArrayRef<NamedAttribute> attrs = {});
Operation* makeAppropriatelyTypedMulOp(OpBuilder& builder, Location loc,
                                       Value lhs, Value rhs,
                                       ArrayRef<NamedAttribute> attrs = {});

inline std::pair<bool, bool> allZeroAllOne(ArrayRef<double> mask) {
  bool allZero = true;
  bool allOne = true;
  for (int64_t v : mask) {
    if (v != 0.0) {
      allZero = false;
    }
    if (v != 1.0) {
      allOne = false;
    }
  }
  return {allZero, allOne};
}

/// Returns true if and only if all types in the given values are the same.
bool allTypesMatch(ArrayRef<Value> values);

/// Extends all (integer-typed) values to the largest bitwidth among them and
/// returns the new list of values.
SmallVector<Value> extendToCommonWidth(
    OpBuilder& b, ArrayRef<Value> values,
    const std::function<void(Operation*)>& createdOpCallback = nullptr);

/// Ensures only one of a Value or Attribute are present
LogicalResult containsExactlyOneOrEmitError(Operation* op, Value dynamicValue,
                                            Attribute staticAttr);
LogicalResult containsExactlyOneOrEmitError(
    Operation* op, Value dynamicValue, std::optional<Attribute> staticAttr);

}  // namespace heir
}  // namespace mlir

#endif  // LIB_UTILS_UTILS_H_
