#pragma once

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/LogicalResult.h"

#include "llvm/ADT/FunctionExtras.h"

namespace onnx_mlir {

using IndexValueResolver = llvm::function_ref<llvm::FailureOr<int64_t>(mlir::Value)>;

mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                    mlir::Location loc,
                                    mlir::AffineMap map,
                                    mlir::ValueRange operands,
                                    mlir::Operation* constantAnchor);

mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                    mlir::Location loc,
                                    mlir::AffineExpr expr,
                                    mlir::ValueRange dims,
                                    mlir::Operation* constantAnchor);

mlir::Value affineMulConst(mlir::OpBuilder& builder,
                           mlir::Location loc,
                           mlir::Value value,
                           int64_t multiplier,
                           mlir::Operation* constantAnchor);

mlir::Value affineAddConst(mlir::OpBuilder& builder,
                           mlir::Location loc,
                           mlir::Value value,
                           int64_t offset,
                           mlir::Operation* constantAnchor);

mlir::Value affineModConst(mlir::OpBuilder& builder,
                           mlir::Location loc,
                           mlir::Value value,
                           int64_t divisor,
                           mlir::Operation* constantAnchor);

mlir::Value affineFloorDivConst(mlir::OpBuilder& builder,
                                mlir::Location loc,
                                mlir::Value value,
                                int64_t divisor,
                                mlir::Operation* constantAnchor);

mlir::Value affineAddModConst(mlir::OpBuilder& builder,
                              mlir::Location loc,
                              mlir::Value value,
                              int64_t offset,
                              int64_t divisor,
                              mlir::Operation* constantAnchor);

mlir::Value affineAddFloorDivConst(mlir::OpBuilder& builder,
                                   mlir::Location loc,
                                   mlir::Value value,
                                   int64_t offset,
                                   int64_t divisor,
                                   mlir::Operation* constantAnchor);

llvm::FailureOr<int64_t>
evaluateAffineExpr(mlir::AffineExpr expr, llvm::ArrayRef<int64_t> dims, llvm::ArrayRef<int64_t> symbols = {});

llvm::FailureOr<int64_t> evaluateSingleResultAffineMap(mlir::AffineMap map, llvm::ArrayRef<int64_t> operands);

llvm::FailureOr<int64_t> evaluateAffineApply(mlir::affine::AffineApplyOp affineApply, IndexValueResolver resolver);

bool isSingleResultSymbolFreeAffineMap(mlir::AffineMap map);

bool isDimAndConstantAffineExpr(mlir::AffineExpr expr);

} // namespace onnx_mlir