#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"

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

using namespace mlir;

namespace onnx_mlir {
namespace {

static bool isSubviewContiguous(const StaticSubviewInfo& info) {
  if (llvm::any_of(info.strides, [](int64_t stride) { return stride != 1; }))
    return false;

  auto sizesAndShape = llvm::zip_equal(llvm::make_range(info.sizes.rbegin(), info.sizes.rend()),
                                       llvm::make_range(info.sourceShape.rbegin(), info.sourceShape.rend()));
  auto firstDifferentSize = std::find_if(sizesAndShape.begin(), sizesAndShape.end(), [&](auto sizeAndShape) -> bool {
    auto [size, dimension] = sizeAndShape;
    return size != dimension;
  });
  if (firstDifferentSize == sizesAndShape.end())
    return true;

  ++firstDifferentSize;
  return std::all_of(firstDifferentSize, sizesAndShape.end(), [](auto sizeAndShape) {
    auto [size, _dimension] = sizeAndShape;
    return size == 1;
  });
}

static OpFoldResult addConstantOffset(OpFoldResult baseOffset, int64_t extraOffset, PatternRewriter& rewriter) {
  if (extraOffset == 0)
    return baseOffset;

  if (auto attr = dyn_cast<Attribute>(baseOffset)) {
    auto integerAttr = dyn_cast<IntegerAttr>(attr);
    assert(integerAttr && "expected integer offset attribute");
    return rewriter.getIndexAttr(integerAttr.getInt() + extraOffset);
  }

  auto value = cast<Value>(baseOffset);
  auto cst = arith::ConstantIndexOp::create(rewriter, value.getLoc(), extraOffset);
  return arith::AddIOp::create(rewriter, value.getLoc(), value, cst).getResult();
}

static Value buildSubviewChunk(const StaticSubviewInfo& info,
                               ArrayRef<int64_t> outerIndices,
                               Location loc,
                               PatternRewriter& rewriter) {
  SmallVector<OpFoldResult> chunkOffsets;
  SmallVector<OpFoldResult> chunkSizes;
  SmallVector<OpFoldResult> chunkStrides;
  chunkOffsets.reserve(info.offsets.size());
  chunkSizes.reserve(info.sizes.size());
  chunkStrides.reserve(info.strides.size());

  for (size_t dim = 0; dim < info.sizes.size(); ++dim) {
    int64_t extraOffset = dim + 1 < info.sizes.size() ? outerIndices[dim] * info.strides[dim] : 0;
    chunkOffsets.push_back(addConstantOffset(info.offsets[dim], extraOffset, rewriter));
    chunkSizes.push_back(rewriter.getIndexAttr(dim + 1 < info.sizes.size() ? 1 : info.sizes.back()));
    chunkStrides.push_back(rewriter.getIndexAttr(info.strides[dim]));
  }

  return memref::SubViewOp::create(rewriter, loc, info.source, chunkOffsets, chunkSizes, chunkStrides);
}

static SmallVector<Value> delinearizeIndexValue(Value linearIndex,
                                                ArrayRef<int64_t> shape,
                                                ArrayRef<int64_t> strides,
                                                PatternRewriter& rewriter) {
  SmallVector<Value> indices;
  indices.reserve(shape.size());

  Value remaining = linearIndex;
  for (auto [_dim, stride] : llvm::enumerate(strides)) {
    auto cStride = arith::ConstantIndexOp::create(rewriter, linearIndex.getLoc(), stride);
    Value index = arith::DivUIOp::create(rewriter, linearIndex.getLoc(), remaining, cStride);
    indices.push_back(index);
    remaining = arith::RemUIOp::create(rewriter, linearIndex.getLoc(), remaining, cStride);
  }

  return indices;
}

static OpFoldResult addDynamicOffset(OpFoldResult baseOffset, Value extraOffset, PatternRewriter& rewriter) {
  if (auto attr = dyn_cast<Attribute>(baseOffset)) {
    auto integerAttr = cast<IntegerAttr>(attr);
    if (integerAttr.getInt() == 0)
      return extraOffset;

    auto cst = arith::ConstantIndexOp::create(rewriter, extraOffset.getLoc(), integerAttr.getInt());
    return arith::AddIOp::create(rewriter, extraOffset.getLoc(), cst, extraOffset).getResult();
  }

  auto value = cast<Value>(baseOffset);
  return arith::AddIOp::create(rewriter, value.getLoc(), value, extraOffset).getResult();
}

static Value buildDynamicSubviewChunk(const StaticSubviewInfo& info,
                                      ArrayRef<Value> outerIndices,
                                      Location loc,
                                      PatternRewriter& rewriter) {
  SmallVector<OpFoldResult> chunkOffsets;
  SmallVector<OpFoldResult> chunkSizes;
  SmallVector<OpFoldResult> chunkStrides;
  chunkOffsets.reserve(info.offsets.size());
  chunkSizes.reserve(info.sizes.size());
  chunkStrides.reserve(info.strides.size());

  for (size_t dim = 0; dim < info.sizes.size(); ++dim) {
    if (dim + 1 < info.sizes.size()) {
      assert(info.strides[dim] == 1 && "loop-based subview rewrite requires unit strides");
      chunkOffsets.push_back(addDynamicOffset(info.offsets[dim], outerIndices[dim], rewriter));
      chunkSizes.push_back(rewriter.getIndexAttr(1));
    }
    else {
      chunkOffsets.push_back(info.offsets[dim]);
      chunkSizes.push_back(rewriter.getIndexAttr(info.sizes.back()));
    }
    chunkStrides.push_back(rewriter.getIndexAttr(info.strides[dim]));
  }

  return memref::SubViewOp::create(rewriter, loc, info.source, chunkOffsets, chunkSizes, chunkStrides);
}

static Value buildContiguousChunk(
  Value source, ArrayRef<int64_t> copyShape, ArrayRef<Value> outerIndices, Location loc, PatternRewriter& rewriter) {
  SmallVector<OpFoldResult> chunkOffsets;
  SmallVector<OpFoldResult> chunkSizes;
  SmallVector<OpFoldResult> chunkStrides;
  chunkOffsets.reserve(copyShape.size());
  chunkSizes.reserve(copyShape.size());
  chunkStrides.reserve(copyShape.size());

  for (size_t dim = 0; dim < copyShape.size(); ++dim) {
    chunkOffsets.push_back(dim + 1 < copyShape.size() ? OpFoldResult(outerIndices[dim]) : rewriter.getIndexAttr(0));
    chunkSizes.push_back(rewriter.getIndexAttr(dim + 1 < copyShape.size() ? 1 : copyShape.back()));
    chunkStrides.push_back(rewriter.getIndexAttr(1));
  }

  return memref::SubViewOp::create(rewriter, loc, source, chunkOffsets, chunkSizes, chunkStrides);
}

template <typename CopyOp, typename CreateCopyOp>
static LogicalResult rewriteSubviewCopyLikeOp(CopyOp copyOp,
                                              Value dst,
                                              Value src,
                                              int64_t dstOffset,
                                              int64_t srcOffset,
                                              int64_t size,
                                              bool allowLoopRewrite,
                                              PatternRewriter& rewriter,
                                              CreateCopyOp createCopyOp) {
  auto srcSubview = getStaticSubviewInfo(src);
  auto dstSubview = getStaticSubviewInfo(dst);
  const bool splitSrc = succeeded(srcSubview) && !isSubviewContiguous(*srcSubview);
  const bool splitDst = succeeded(dstSubview) && !isSubviewContiguous(*dstSubview);
  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 && (srcOffset != 0 || llvm::any_of(srcSubview->strides, [](int64_t stride) { return stride != 1; })))
    return failure();
  if (splitDst && (dstOffset != 0 || 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);

  if (allowLoopRewrite && numSlices > 1 && srcOffset == 0 && dstOffset == 0
      && sourceType.getRank() == static_cast<int64_t>(copyShape.size())
      && dstType.getRank() == static_cast<int64_t>(copyShape.size())) {
    auto c0 = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), 0);
    auto cUpper = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), numSlices);
    auto cStep = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), 1);

    auto loop = scf::ForOp::create(rewriter, copyOp.getLoc(), c0, cUpper, cStep, ValueRange {});
    rewriter.setInsertionPointToStart(loop.getBody());

    SmallVector<Value> outerIndices =
      outerShape.empty() ? SmallVector<Value> {}
                         : delinearizeIndexValue(loop.getInductionVar(), outerShape, outerStrides, rewriter);
    Value chunkDst = splitDst ? buildDynamicSubviewChunk(*dstSubview, outerIndices, copyOp.getLoc(), rewriter)
                              : buildContiguousChunk(dst, copyShape, outerIndices, copyOp.getLoc(), rewriter);
    Value chunkSrc = splitSrc ? buildDynamicSubviewChunk(*srcSubview, outerIndices, copyOp.getLoc(), rewriter)
                              : buildContiguousChunk(src, copyShape, outerIndices, copyOp.getLoc(), rewriter);
    createCopyOp(cast<MemRefType>(chunkDst.getType()), chunkDst, chunkSrc, 0, 0, sliceBytes);
    return success();
  }

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

  return success();
}

// Splits core copies through subviews into contiguous copy chunks for codegen.
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(),
      /*allowLoopRewrite=*/true,
      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();
  }
};

// Splits host-to-device subview loads into contiguous copy chunks.
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(),
      /*allowLoopRewrite=*/true,
      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();
  }
};

// Splits device-to-host subview stores into contiguous copy chunks.
struct RewriteHostSubviewStorePattern final : OpRewritePattern<pim::PimMemCopyDevToHostOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(pim::PimMemCopyDevToHostOp copyOp, PatternRewriter& rewriter) const override {
    auto status = rewriteSubviewCopyLikeOp(
      copyOp,
      copyOp.getHostTarget(),
      copyOp.getDeviceSource(),
      copyOp.getHostTargetOffset(),
      copyOp.getDeviceSourceOffset(),
      copyOp.getSize(),
      /*allowLoopRewrite=*/false,
      rewriter,
      [&](
        MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
        pim::PimMemCopyDevToHostOp::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.getHostTarget());
    return success();
  }
};

// Folds constant subviews used as core weights into standalone globals.
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 staticOffsets = getStaticSubviewOffsets(*subviewInfo);
    if (failed(staticOffsets))
      return failure();

    auto elementType = cast<MemRefType>(subviewOp.getType()).getElementType();
    auto resultMemRefType =
      MemRefType::get(SmallVector<int64_t>(subviewInfo->sizes.begin(), subviewInfo->sizes.end()), elementType);
    auto foldedAttr = foldDenseSubview(*denseAttr, *staticOffsets, resultMemRefType.getShape());
    if (failed(foldedAttr))
      return failure();

    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,
               RewriteHostSubviewStorePattern,
               FoldConstantCoreSubviewPattern>(patterns.getContext());
}

} // namespace onnx_mlir