#pragma once

#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Transforms/DialectConversion.h"

#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"

namespace onnx_mlir {

struct MixedSliceGeometry {
  llvm::SmallVector<mlir::OpFoldResult> offsets;
  llvm::SmallVector<mlir::OpFoldResult> sizes;
  llvm::SmallVector<mlir::OpFoldResult> strides;
};

mlir::Value extractAxisSlice(
  mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);

mlir::Value extractStaticSliceOrIdentity(mlir::OpBuilder& rewriter,
                                         mlir::Location loc,
                                         mlir::Value source,
                                         mlir::RankedTensorType resultType,
                                         llvm::ArrayRef<mlir::OpFoldResult> offsets,
                                         llvm::ArrayRef<mlir::OpFoldResult> sizes,
                                         llvm::ArrayRef<mlir::OpFoldResult> strides);

mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
                              mlir::Location loc,
                              mlir::Value source,
                              mlir::Value dest,
                              llvm::ArrayRef<mlir::OpFoldResult> offsets);

mlir::Value extractMixedSliceOrIdentity(mlir::OpBuilder &rewriter,
                                        mlir::Location loc,
                                        mlir::Value source,
                                        mlir::RankedTensorType resultType,
                                        const MixedSliceGeometry &geometry);

mlir::Value insertMixedSlice(mlir::OpBuilder &builder,
                             mlir::Location loc,
                             mlir::Value source,
                             mlir::Value dest,
                             const MixedSliceGeometry &geometry);

mlir::FailureOr<mlir::Value>
addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value);

mlir::FailureOr<mlir::Value> removeLeadingUnitTensorDimension(
  mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, mlir::RankedTensorType resultType);

} // namespace onnx_mlir