#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

struct PackSpatialConcatInputsPattern final : OpRewritePattern<spatial::SpatConcatOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatConcatOp concatOp, PatternRewriter& rewriter) const override {
    if (concatOp.getAxis() != 0 || concatOp.getInputs().empty())
      return failure();

    SmallVector<Value> packedInputs;
    bool changed = false;

    for (unsigned index = 0; index < concatOp.getInputs().size();) {
      Value input = concatOp.getInputs()[index];

      if (input.getDefiningOp<tensor::ExtractSliceOp>()) {
        unsigned endIndex = index + 1;
        while (endIndex < concatOp.getInputs().size()
               && concatOp.getInputs()[endIndex].getDefiningOp<tensor::ExtractSliceOp>())
          ++endIndex;

        Value packedInput = createPackedExtractSliceTensor(
          concatOp.getInputs().slice(index, endIndex - index), rewriter, concatOp.getLoc());
        if (packedInput) {
          packedInputs.push_back(packedInput);
          changed = true;
          index = endIndex;
          continue;
        }
      }

      auto result = dyn_cast<OpResult>(input);
      if (!result) {
        packedInputs.push_back(input);
        ++index;
        continue;
      }

      Operation* owner = result.getOwner();
      unsigned startIndex = result.getResultNumber();
      unsigned endIndex = index + 1;
      while (endIndex < concatOp.getInputs().size()) {
        auto nextResult = dyn_cast<OpResult>(concatOp.getInputs()[endIndex]);
        if (!nextResult || nextResult.getOwner() != owner
            || nextResult.getResultNumber() != startIndex + (endIndex - index))
          break;
        ++endIndex;
      }

      unsigned count = endIndex - index;
      Value packedInput;
      if (auto extractRowsOp = dyn_cast<spatial::SpatExtractRowsOp>(owner))
        packedInput = createPackedExtractRowsSlice(extractRowsOp, startIndex, count, rewriter, concatOp.getLoc());

      if (packedInput) {
        packedInputs.push_back(packedInput);
        changed = true;
      }
      else {
        for (unsigned oldIndex = index; oldIndex < endIndex; ++oldIndex)
          packedInputs.push_back(concatOp.getInputs()[oldIndex]);
      }

      index = endIndex;
    }

    if (!changed)
      return failure();

    auto outputType = cast<ShapedType>(concatOp.getOutput().getType());
    auto newConcat = pim::PimConcatOp::create(
      rewriter,
      concatOp.getLoc(),
      concatOp.getOutput().getType(),
      concatOp.getAxisAttr(),
      ValueRange(packedInputs),
      tensor::EmptyOp::create(rewriter, concatOp.getLoc(), outputType.getShape(), outputType.getElementType())
        .getResult());
    rewriter.replaceOp(concatOp, newConcat.getOutput());
    return success();
  }
};

} // namespace

RankedTensorType getPackedTensorType(RankedTensorType elementType, int64_t count) {
  SmallVector<int64_t> packedShape(elementType.getShape().begin(), elementType.getShape().end());
  packedShape[0] *= count;
  return RankedTensorType::get(packedShape, elementType.getElementType());
}

Value extractPackedChunk(
  Value packedValue, RankedTensorType chunkType, unsigned index, OpBuilder& builder, Location loc) {
  auto packedType = dyn_cast<RankedTensorType>(packedValue.getType());
  if (packedType && packedType == chunkType && index == 0)
    return packedValue;

  SmallVector<OpFoldResult> offsets;
  SmallVector<OpFoldResult> sizes;
  SmallVector<OpFoldResult> strides;
  offsets.reserve(chunkType.getRank());
  sizes.reserve(chunkType.getRank());
  strides.reserve(chunkType.getRank());

  offsets.push_back(builder.getIndexAttr(static_cast<int64_t>(index) * chunkType.getDimSize(0)));
  sizes.push_back(builder.getIndexAttr(chunkType.getDimSize(0)));
  strides.push_back(builder.getIndexAttr(1));
  for (int64_t dim = 1; dim < chunkType.getRank(); ++dim) {
    offsets.push_back(builder.getIndexAttr(0));
    sizes.push_back(builder.getIndexAttr(chunkType.getDimSize(dim)));
    strides.push_back(builder.getIndexAttr(1));
  }

  return tensor::ExtractSliceOp::create(builder, loc, chunkType, packedValue, offsets, sizes, strides).getResult();
}

Value createPackedExtractRowsSlice(
  spatial::SpatExtractRowsOp extractRowsOp, unsigned startIndex, unsigned count, OpBuilder& builder, Location loc) {
  auto rowType = dyn_cast<RankedTensorType>(extractRowsOp.getOutputs()[startIndex].getType());
  auto inputType = dyn_cast<RankedTensorType>(extractRowsOp.getInput().getType());
  if (!rowType || !inputType || !rowType.hasStaticShape() || !inputType.hasStaticShape() || rowType.getRank() == 0)
    return {};

  int64_t rowsPerValue = rowType.getDimSize(0);
  if (ShapedType::isDynamic(rowsPerValue))
    return {};

  auto packedType = getPackedTensorType(rowType, static_cast<int64_t>(count));
  SmallVector<OpFoldResult> offsets;
  SmallVector<OpFoldResult> sizes;
  SmallVector<OpFoldResult> strides;
  offsets.reserve(inputType.getRank());
  sizes.reserve(inputType.getRank());
  strides.reserve(inputType.getRank());

  offsets.push_back(builder.getIndexAttr(static_cast<int64_t>(startIndex) * rowsPerValue));
  sizes.push_back(builder.getIndexAttr(static_cast<int64_t>(count) * rowsPerValue));
  strides.push_back(builder.getIndexAttr(1));
  for (int64_t dim = 1; dim < inputType.getRank(); ++dim) {
    offsets.push_back(builder.getIndexAttr(0));
    sizes.push_back(builder.getIndexAttr(inputType.getDimSize(dim)));
    strides.push_back(builder.getIndexAttr(1));
  }

  bool coversWholeSource = packedType == inputType && startIndex == 0;
  if (coversWholeSource)
    return extractRowsOp.getInput();

  return tensor::ExtractSliceOp::create(builder, loc, packedType, extractRowsOp.getInput(), offsets, sizes, strides)
    .getResult();
}

Value createPackedExtractSliceTensor(ValueRange values, OpBuilder& builder, Location loc) {
  if (values.empty())
    return {};
  if (values.size() == 1)
    return values.front();

  auto firstSliceOp = values.front().getDefiningOp<tensor::ExtractSliceOp>();
  if (!firstSliceOp)
    return {};

  auto firstType = dyn_cast<RankedTensorType>(firstSliceOp.getResult().getType());
  auto sourceType = dyn_cast<RankedTensorType>(firstSliceOp.getSource().getType());
  if (!firstType || !sourceType || !firstType.hasStaticShape() || !sourceType.hasStaticShape()
      || firstType.getRank() == 0)
    return {};

  auto hasStaticValues = [](ArrayRef<int64_t> values) {
    return llvm::all_of(values, [](int64_t value) { return !ShapedType::isDynamic(value); });
  };
  if (!hasStaticValues(firstSliceOp.getStaticOffsets()) || !hasStaticValues(firstSliceOp.getStaticSizes())
      || !hasStaticValues(firstSliceOp.getStaticStrides()))
    return {};

  ArrayRef<int64_t> firstOffsets = firstSliceOp.getStaticOffsets();
  ArrayRef<int64_t> firstSizes = firstSliceOp.getStaticSizes();
  ArrayRef<int64_t> firstStrides = firstSliceOp.getStaticStrides();
  int64_t rowsPerValue = firstSizes[0];
  if (ShapedType::isDynamic(rowsPerValue))
    return {};

  for (size_t index = 1; index < values.size(); ++index) {
    auto sliceOp = values[index].getDefiningOp<tensor::ExtractSliceOp>();
    if (!sliceOp || sliceOp.getSource() != firstSliceOp.getSource()
        || sliceOp.getResult().getType() != firstSliceOp.getResult().getType()
        || !hasStaticValues(sliceOp.getStaticOffsets()) || !hasStaticValues(sliceOp.getStaticSizes())
        || !hasStaticValues(sliceOp.getStaticStrides()))
      return {};

    if (sliceOp.getStaticSizes() != firstSizes || sliceOp.getStaticStrides() != firstStrides)
      return {};

    if (sliceOp.getStaticOffsets()[0] != firstOffsets[0] + static_cast<int64_t>(index) * rowsPerValue)
      return {};

    for (int64_t dim = 1; dim < firstType.getRank(); ++dim)
      if (sliceOp.getStaticOffsets()[dim] != firstOffsets[dim])
        return {};
  }

  auto packedType = getPackedTensorType(firstType, static_cast<int64_t>(values.size()));
  SmallVector<OpFoldResult> offsets;
  SmallVector<OpFoldResult> sizes;
  SmallVector<OpFoldResult> strides;
  offsets.reserve(firstType.getRank());
  sizes.reserve(firstType.getRank());
  strides.reserve(firstType.getRank());

  offsets.push_back(builder.getIndexAttr(firstOffsets[0]));
  sizes.push_back(builder.getIndexAttr(rowsPerValue * static_cast<int64_t>(values.size())));
  strides.push_back(builder.getIndexAttr(firstStrides[0]));
  for (int64_t dim = 1; dim < firstType.getRank(); ++dim) {
    offsets.push_back(builder.getIndexAttr(firstOffsets[dim]));
    sizes.push_back(builder.getIndexAttr(firstSizes[dim]));
    strides.push_back(builder.getIndexAttr(firstStrides[dim]));
  }

  bool coversWholeSource = packedType == sourceType;
  for (int64_t dim = 0; coversWholeSource && dim < sourceType.getRank(); ++dim)
    coversWholeSource = firstOffsets[dim] == 0 && firstStrides[dim] == 1;
  if (coversWholeSource)
    return firstSliceOp.getSource();

  return tensor::ExtractSliceOp::create(builder, loc, packedType, firstSliceOp.getSource(), offsets, sizes, strides)
    .getResult();
}

void populateTensorPackingPatterns(RewritePatternSet& patterns) {
  patterns.add<PackSpatialConcatInputsPattern>(patterns.getContext());
}

void eraseUnusedTensorPackingOps(func::FuncOp funcOp, IRRewriter& rewriter) {
  auto eraseUnusedOps = [&](auto tag) {
    using OpTy = decltype(tag);
    SmallVector<OpTy> ops;
    funcOp.walk([&](OpTy op) { ops.push_back(op); });
    for (auto op : llvm::reverse(ops))
      if (op->use_empty())
        rewriter.eraseOp(op);
  };
  eraseUnusedOps(tensor::ConcatOp {});
  eraseUnusedOps(tensor::ExtractSliceOp {});
  eraseUnusedOps(spatial::SpatExtractRowsOp {});
}

} // namespace onnx_mlir