#include "../Common.hpp"
#include "../Patterns.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

template <typename CopyOp, typename CreateCopyOp>
static LogicalResult rewriteSubviewCopyLikeOp(CopyOp copyOp,
                                              Value dst,
                                              Value src,
                                              int64_t dstOffset,
                                              int64_t srcOffset,
                                              int64_t size,
                                              PatternRewriter& rewriter,
                                              CreateCopyOp createCopyOp) {
  auto srcSubview = getStaticSubviewInfo(src);
  auto dstSubview = getStaticSubviewInfo(dst);
  const bool splitSrc =
    succeeded(srcSubview)
    && !isMemoryContiguous(srcSubview->sourceShape, srcSubview->offsets, srcSubview->sizes, srcSubview->strides);
  const bool splitDst =
    succeeded(dstSubview)
    && !isMemoryContiguous(dstSubview->sourceShape, dstSubview->offsets, dstSubview->sizes, dstSubview->strides);
  if (!splitSrc && !splitDst)
    return failure();

  auto sourceType = dyn_cast<MemRefType>(src.getType());
  auto dstType = dyn_cast<MemRefType>(dst.getType());
  if (!sourceType || !dstType || !sourceType.hasStaticShape() || !dstType.hasStaticShape())
    return failure();
  if (sourceType.getElementType() != dstType.getElementType())
    return failure();

  if (splitSrc && llvm::any_of(srcSubview->strides, [](int64_t stride) { return stride != 1; }))
    return failure();
  if (splitDst && llvm::any_of(dstSubview->strides, [](int64_t stride) { return stride != 1; }))
    return failure();

  ArrayRef<int64_t> copyShape = splitSrc ? ArrayRef<int64_t>(srcSubview->sizes) : ArrayRef<int64_t>(dstSubview->sizes);
  if (splitSrc && splitDst && copyShape != ArrayRef<int64_t>(dstSubview->sizes))
    return failure();

  const int64_t elementByteWidth = sourceType.getElementTypeBitWidth() / 8;
  if (elementByteWidth <= 0)
    return failure();

  const int64_t totalBytes = getNumElements(copyShape) * elementByteWidth;
  if (size != totalBytes)
    return failure();

  const int64_t sliceBytes = copyShape.back() * elementByteWidth;
  if (sliceBytes <= 0)
    return failure();

  SmallVector<int64_t> outerShape(copyShape.begin(), copyShape.end() - 1);
  auto outerStrides = computeRowMajorStrides(outerShape);
  const int64_t numSlices = outerShape.empty() ? 1 : getNumElements(outerShape);

  rewriter.setInsertionPoint(copyOp);
  for (int64_t linearIndex = 0; linearIndex < numSlices; ++linearIndex) {
    SmallVector<int64_t> outerIndices =
      outerShape.empty() ? SmallVector<int64_t> {} : delinearizeIndex(linearIndex, outerShape, outerStrides);
    const int64_t srcByteOffset =
      srcOffset
      + (splitSrc ? getSubviewChunkOffsetBytes(*srcSubview, outerIndices, elementByteWidth) : linearIndex * sliceBytes);
    const int64_t dstByteOffset =
      dstOffset
      + (splitDst ? getSubviewChunkOffsetBytes(*dstSubview, outerIndices, elementByteWidth) : linearIndex * sliceBytes);
    createCopyOp(splitDst ? cast<MemRefType>(dstSubview->source.getType()) : dstType,
                 splitDst ? dstSubview->source : dst,
                 splitSrc ? srcSubview->source : src,
                 dstByteOffset,
                 srcByteOffset,
                 sliceBytes);
  }

  return success();
}

struct RewriteCoreSubviewCopyPattern final : OpRewritePattern<pim::PimMemCopyOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(pim::PimMemCopyOp copyOp, PatternRewriter& rewriter) const override {
    if (!copyOp->getParentOfType<pim::PimCoreOp>())
      return failure();

    auto status = rewriteSubviewCopyLikeOp(
      copyOp,
      copyOp.getTarget(),
      copyOp.getSource(),
      copyOp.getTargetOffset(),
      copyOp.getSourceOffset(),
      copyOp.getSize(),
      rewriter,
      [&](
        MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
        pim::PimMemCopyOp::create(rewriter,
                                  copyOp.getLoc(),
                                  resultType,
                                  dst,
                                  src,
                                  rewriter.getI32IntegerAttr(static_cast<int32_t>(dstByteOffset)),
                                  rewriter.getI32IntegerAttr(static_cast<int32_t>(srcByteOffset)),
                                  rewriter.getI32IntegerAttr(static_cast<int32_t>(sliceBytes)));
      });
    if (failed(status))
      return failure();

    rewriter.replaceOp(copyOp, copyOp.getTarget());
    return success();
  }
};

struct RewriteHostSubviewLoadPattern final : OpRewritePattern<pim::PimMemCopyHostToDevOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(pim::PimMemCopyHostToDevOp copyOp, PatternRewriter& rewriter) const override {
    auto status = rewriteSubviewCopyLikeOp(
      copyOp,
      copyOp.getDeviceTarget(),
      copyOp.getHostSource(),
      copyOp.getDeviceTargetOffset(),
      copyOp.getHostSourceOffset(),
      copyOp.getSize(),
      rewriter,
      [&](
        MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
        pim::PimMemCopyHostToDevOp::create(rewriter,
                                           copyOp.getLoc(),
                                           resultType,
                                           dst,
                                           src,
                                           rewriter.getI32IntegerAttr(static_cast<int32_t>(dstByteOffset)),
                                           rewriter.getI32IntegerAttr(static_cast<int32_t>(srcByteOffset)),
                                           rewriter.getI32IntegerAttr(static_cast<int32_t>(sliceBytes)));
      });
    if (failed(status))
      return failure();

    rewriter.replaceOp(copyOp, copyOp.getDeviceTarget());
    return success();
  }
};

struct FoldConstantCoreSubviewPattern final : OpRewritePattern<memref::SubViewOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(memref::SubViewOp subviewOp, PatternRewriter& rewriter) const override {
    if (subviewOp.use_empty())
      return failure();
    if (!llvm::all_of(subviewOp->getUsers(), [](Operation* user) { return isa<pim::PimCoreOp>(user); }))
      return failure();

    auto moduleOp = subviewOp->getParentOfType<ModuleOp>();
    if (!moduleOp)
      return failure();
    auto denseAttr = getDenseGlobalValue(moduleOp, stripMemRefCasts(subviewOp.getSource()));
    if (failed(denseAttr))
      return failure();

    auto subviewInfo = getStaticSubviewInfo(subviewOp.getResult());
    if (failed(subviewInfo))
      return failure();
    if (llvm::any_of(subviewInfo->strides, [](int64_t stride) { return stride != 1; }))
      return failure();

    auto sourceType = dyn_cast<RankedTensorType>(denseAttr->getType());
    if (!sourceType || !sourceType.hasStaticShape())
      return failure();

    auto elementType = cast<MemRefType>(subviewOp.getType()).getElementType();
    auto resultMemRefType =
      MemRefType::get(SmallVector<int64_t>(subviewInfo->sizes.begin(), subviewInfo->sizes.end()), elementType);
    auto resultTensorType = RankedTensorType::get(resultMemRefType.getShape(), elementType);
    const int64_t numResultElements = resultTensorType.getNumElements();

    auto sourceStrides = computeRowMajorStrides(sourceType.getShape());
    auto resultStrides = computeRowMajorStrides(resultTensorType.getShape());
    SmallVector<Attribute> sourceValues(denseAttr->getValues<Attribute>());
    SmallVector<Attribute> resultValues(numResultElements);
    for (int64_t i = 0; i < numResultElements; ++i) {
      auto resultIndices = delinearizeIndex(i, resultTensorType.getShape(), resultStrides);
      SmallVector<int64_t> sourceIndices;
      sourceIndices.reserve(resultIndices.size());
      for (auto [off, idx] : llvm::zip_equal(subviewInfo->offsets, resultIndices))
        sourceIndices.push_back(off + idx);
      resultValues[i] = sourceValues[linearizeIndex(sourceIndices, sourceStrides)];
    }
    auto foldedAttr = DenseElementsAttr::get(resultTensorType, resultValues);

    auto newGlobal =
      createFoldedGlobal(moduleOp, subviewOp.getLoc(), resultMemRefType, foldedAttr, "pim_folded_subview");
    markWeightAlways(newGlobal);

    rewriter.setInsertionPoint(subviewOp);
    auto newGetGlobal =
      memref::GetGlobalOp::create(rewriter, subviewOp.getLoc(), resultMemRefType, newGlobal.getName());
    markWeightAlways(newGetGlobal);

    rewriter.replaceOp(subviewOp, newGetGlobal.getResult());
    return success();
  }
};

} // namespace

void populateConstantFoldingSubviewPatterns(RewritePatternSet& patterns) {
  patterns.add<RewriteCoreSubviewCopyPattern, RewriteHostSubviewLoadPattern, FoldConstantCoreSubviewPattern>(
    patterns.getContext());
}

} // namespace onnx_mlir
