#include "mlir/Dialect/Tensor/IR/Tensor.h"

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

using namespace mlir;

namespace onnx_mlir {

Value extractAxisSlice(
  PatternRewriter& rewriter, Location loc, Value source, int64_t axis, int64_t offset, int64_t size) {
  auto sourceType = cast<RankedTensorType>(source.getType());
  SmallVector<int64_t> resultShape(sourceType.getShape());
  resultShape[axis] = size;
  auto resultType = RankedTensorType::get(resultShape, sourceType.getElementType(), sourceType.getEncoding());

  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, sourceType.getRank());
  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, sourceType.getShape());
  offsets[axis] = rewriter.getIndexAttr(offset);
  sizes[axis] = rewriter.getIndexAttr(size);
  return tensor::ExtractSliceOp::create(
           rewriter, loc, resultType, source, offsets, sizes, getUnitStrides(rewriter, sourceType.getRank()))
    .getResult();
}

Value extractStaticSliceOrIdentity(RewriterBase& rewriter,
                                   Location loc,
                                   Value source,
                                   RankedTensorType resultType,
                                   ArrayRef<OpFoldResult> offsets,
                                   ArrayRef<OpFoldResult> sizes,
                                   ArrayRef<OpFoldResult> strides) {
  auto sourceType = cast<RankedTensorType>(source.getType());
  size_t rank = static_cast<size_t>(sourceType.getRank());

  bool isIdentitySlice =
    sourceType == resultType && sourceType.hasStaticShape() && offsets.size() == rank && sizes.size() == rank
    && strides.size() == rank;
  if (isIdentitySlice) {
    ArrayRef<int64_t> sourceShape = sourceType.getShape();
    for (auto [dim, offset, size, stride] : llvm::zip_equal(sourceShape, offsets, sizes, strides)) {
      std::optional<int64_t> staticOffset = mlir::getConstantIntValue(offset);
      std::optional<int64_t> staticSize = mlir::getConstantIntValue(size);
      std::optional<int64_t> staticStride = mlir::getConstantIntValue(stride);
      if (!staticOffset || !staticSize || !staticStride || *staticOffset != 0 || *staticSize != dim
          || *staticStride != 1) {
        isIdentitySlice = false;
        break;
      }
    }
  }

  if (isIdentitySlice)
    return source;

  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, strides).getResult();
}

Value insertStaticSlice(
  PatternRewriter& rewriter, Location loc, Value source, Value dest, ArrayRef<OpFoldResult> offsets) {
  auto sourceType = cast<RankedTensorType>(source.getType());
  return tensor::InsertSliceOp::create(rewriter,
                                       loc,
                                       source,
                                       dest,
                                       offsets,
                                       getStaticSizes(rewriter, sourceType.getShape()),
                                       getUnitStrides(rewriter, sourceType.getRank()))
    .getResult();
}

Value extractMixedSliceOrIdentity(RewriterBase &rewriter,
                                  Location loc,
                                  Value source,
                                  RankedTensorType resultType,
                                  const MixedSliceGeometry &geometry) {
  return extractStaticSliceOrIdentity(rewriter, loc, source, resultType,
                                      geometry.offsets, geometry.sizes,
                                      geometry.strides);
}

Value insertMixedSlice(OpBuilder &builder, Location loc, Value source,
                       Value dest, const MixedSliceGeometry &geometry) {
  return tensor::InsertSliceOp::create(builder, loc, source, dest,
                                       geometry.offsets, geometry.sizes,
                                       geometry.strides);
}

FailureOr<Value> addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) {
  auto type = dyn_cast<RankedTensorType>(value.getType());
  if (!type || !type.hasStaticShape())
    return failure();
  SmallVector<int64_t> shape {1};
  llvm::append_range(shape, type.getShape());
  auto resultType = RankedTensorType::get(shape, type.getElementType(), type.getEncoding());
  SmallVector<ReassociationIndices> reassociation;
  if (type.getRank() != 0) {
    reassociation.push_back({0, 1});
    for (int64_t dim = 1; dim < type.getRank(); ++dim)
      reassociation.push_back({dim + 1});
  }
  return tensor::ExpandShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
}

FailureOr<Value> removeLeadingUnitTensorDimension(
  OpBuilder& builder, Location loc, Value value, RankedTensorType resultType) {
  if (value.getType() == resultType)
    return value;
  auto type = dyn_cast<RankedTensorType>(value.getType());
  if (!type || !resultType || !type.hasStaticShape() || !resultType.hasStaticShape()
      || type.getRank() != resultType.getRank() + 1 || type.getDimSize(0) != 1
      || type.getElementType() != resultType.getElementType()
      || !llvm::equal(type.getShape().drop_front(), resultType.getShape()))
    return failure();
  SmallVector<ReassociationIndices> reassociation;
  if (resultType.getRank() != 0) {
    reassociation.push_back({0, 1});
    for (int64_t dim = 1; dim < resultType.getRank(); ++dim)
      reassociation.push_back({dim + 1});
  }
  return tensor::CollapseShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
}

} // namespace onnx_mlir