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

#include "llvm/ADT/SmallVector.h"

#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

static bool haveStaticPositiveShape(ArrayRef<int64_t> shape) {
  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
}

static bool inferCollapseReassociation(ArrayRef<int64_t> sourceShape,
                                       ArrayRef<int64_t> resultShape,
                                       SmallVector<ReassociationIndices>& reassociation) {
  reassociation.clear();

  size_t sourceIdx = 0;
  size_t resultIdx = 0;
  while (sourceIdx < sourceShape.size() && resultIdx < resultShape.size()) {
    int64_t sourceProduct = sourceShape[sourceIdx];
    int64_t resultProduct = resultShape[resultIdx];

    ReassociationIndices group;
    group.push_back(sourceIdx);
    while (sourceProduct != resultProduct) {
      if (sourceProduct > resultProduct)
        return false;
      sourceIdx++;
      if (sourceIdx >= sourceShape.size())
        return false;
      group.push_back(sourceIdx);
      sourceProduct *= sourceShape[sourceIdx];
    }

    reassociation.push_back(group);
    sourceIdx++;
    resultIdx++;
  }

  return sourceIdx == sourceShape.size() && resultIdx == resultShape.size();
}

static bool inferExpandReassociation(ArrayRef<int64_t> sourceShape,
                                     ArrayRef<int64_t> resultShape,
                                     SmallVector<ReassociationIndices>& reassociation) {
  reassociation.clear();

  size_t sourceIdx = 0;
  size_t resultIdx = 0;
  while (sourceIdx < sourceShape.size() && resultIdx < resultShape.size()) {
    int64_t sourceProduct = sourceShape[sourceIdx];
    int64_t resultProduct = resultShape[resultIdx];

    ReassociationIndices group;
    group.push_back(resultIdx);
    while (resultProduct != sourceProduct) {
      if (resultProduct > sourceProduct)
        return false;
      resultIdx++;
      if (resultIdx >= resultShape.size())
        return false;
      group.push_back(resultIdx);
      resultProduct *= resultShape[resultIdx];
    }

    reassociation.push_back(group);
    sourceIdx++;
    resultIdx++;
  }

  return sourceIdx == sourceShape.size() && resultIdx == resultShape.size();
}

struct Reshape : OpConversionPattern<ONNXReshapeOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(ONNXReshapeOp reshapeOp,
                                ONNXReshapeOpAdaptor adaptor,
                                ConversionPatternRewriter& rewriter) const override {
    auto sourceType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
    auto resultType = dyn_cast<RankedTensorType>(reshapeOp.getReshaped().getType());
    if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
      return failure();
    if (!haveStaticPositiveShape(sourceType.getShape()) || !haveStaticPositiveShape(resultType.getShape()))
      return failure();

    if (sourceType == resultType) {
      rewriter.replaceOp(reshapeOp, adaptor.getData());
      return success();
    }

    SmallVector<ReassociationIndices> reassociation;
    if (sourceType.getRank() > resultType.getRank()
        && inferCollapseReassociation(sourceType.getShape(), resultType.getShape(), reassociation)) {
      rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(reshapeOp, resultType, adaptor.getData(), reassociation);
      return success();
    }

    if (sourceType.getRank() < resultType.getRank()
        && inferExpandReassociation(sourceType.getShape(), resultType.getShape(), reassociation)) {
      rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(reshapeOp, resultType, adaptor.getData(), reassociation);
      return success();
    }

    return failure();
  }
};

} // namespace

void populateReshapePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.insert<Reshape>(ctx); }

} // namespace onnx_mlir
