#include "lib/Transforms/ElementwiseToAffine/ElementwiseToAffine.h"

#include <cstddef>
#include <cstdint>
#include <utility>

#include "lib/Dialect/HEIRInterfaces.h"
#include "llvm/include/llvm/ADT/STLExtras.h"    // from @llvm-project
#include "llvm/include/llvm/ADT/SmallVector.h"  // from @llvm-project
#include "llvm/include/llvm/Support/Casting.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
#include "mlir/include/mlir/IR/OpDefinition.h"           // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
#include "mlir/include/mlir/IR/Visitors.h"               // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
#include "mlir/include/mlir/Support/LogicalResult.h"     // from @llvm-project
#include "mlir/include/mlir/Support/WalkResult.h"        // from @llvm-project
#include "mlir/include/mlir/Transforms/DialectConversion.h"  // from @llvm-project

namespace mlir {
namespace heir {

#define GEN_PASS_DEF_ELEMENTWISETOAFFINE
#include "lib/Transforms/ElementwiseToAffine/ElementwiseToAffine.h.inc"

// All of this is based on the ElementwiseToLinalg Pass in
// mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp

// An op is supported if it is ElementwiseMappable (all operands and results
// are mapped over), or if it has the ElementwiseByOperandOpInterface, which
// means only some operands are mapped over.
static bool isSupported(Operation* op) {
  // Skip kernel dialect operations if they operate on standard floats or
  // integers.
  if (op->getDialect()->getNamespace() == "kernel") {
    for (Type type : op->getOperandTypes()) {
      if (auto tensorType = llvm::dyn_cast<RankedTensorType>(type)) {
        Type eltType = tensorType.getElementType();
        if (llvm::isa<FloatType, IntegerType>(eltType)) {
          return false;
        }
      }
    }
  }

  auto hasRankedTensorResult = [](Operation* op) {
    return llvm::any_of(op->getResultTypes(),
                        [](Type type) { return isa<RankedTensorType>(type); });
  };

  if (OpTrait::hasElementwiseMappableTraits(op)) {
    return llvm::any_of(
               op->getOperandTypes(),
               [](Type type) { return isa<RankedTensorType>(type); }) &&
           hasRankedTensorResult(op);
  }

  if (auto interface = dyn_cast<ElementwiseByOperandOpInterface>(op)) {
    if (!hasRankedTensorResult(op)) return false;

    // Ensure all mappable operands are tensors
    if (!llvm::all_of(
            llvm::enumerate(op->getOperands()), [&](auto indexedOperand) {
              return !interface.operandIsMappable(indexedOperand.index()) ||
                     isa<RankedTensorType>(indexedOperand.value().getType());
            }))
      return false;

    // An op with the interface is supported (needs processing) if at least one
    // mappable operand has dimensions to map over.
    return llvm::any_of(llvm::enumerate(op->getOperands()),
                        [&](auto indexedOperand) {
                          return interface.operandIsMappable(indexedOperand.index()) &&
                 !interface
                      .mappedDimensionsForOperand(indexedOperand.index())
                      .empty();
                        });
  }

  return false;
}

namespace {

struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
  ConvertAnyElementwiseMappableOpOnRankedTensors(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
  LogicalResult matchAndRewrite(Operation* op,
                                PatternRewriter& rewriter) const final {
    if (!isSupported(op))
      return rewriter.notifyMatchFailure(
          op,
          "unsupported op: not elementwise mappable, or doesn't have "
          "ElementwiseByOperandOpInterface, or doesn't have tensor operands");

    // Find the tensor result to determine loop nest.
    // We assume there's exactly one tensor result for now, but other results
    // can be non-tensors.
    RankedTensorType resultType;
    int tensorResultIndex = -1;
    for (auto [i, type] : llvm::enumerate(op->getResultTypes())) {
      if (auto rankedTensorTy = dyn_cast<RankedTensorType>(type)) {
        if (tensorResultIndex != -1) {
          return rewriter.notifyMatchFailure(
              op, "unsupported op: more than one tensor result");
        }
        resultType = rankedTensorTy;
        tensorResultIndex = i;
      }
    }

    if (tensorResultIndex == -1) {
      return rewriter.notifyMatchFailure(op,
                                         "unsupported op: no tensor result");
    }

    Type elementType = resultType.getElementType();
    ArrayRef<int64_t> shape = resultType.getShape();
    int64_t rank = resultType.getRank();

    // Save insertion point prior to entering loop nest
    OpBuilder::InsertionGuard guard(rewriter);

    // Create an empty tensor as initial value of the iter_args
    Value target = tensor::EmptyOp::create(
        rewriter, op->getLoc(), shape, elementType, resultType.getEncoding());
    llvm::SmallVector<Value, 1> indices;

    // Create a an affine.for loop nest of depth rank
    for (int64_t i = 0; i < rank; ++i) {
      auto loop =
          affine::AffineForOp::create(rewriter, op->getLoc(), /* lowerBound*/ 0,
                                      /* upperBound*/ shape[i],
                                      /* step*/ 1,
                                      /* iterArgs*/ target);
      target = loop.getRegionIterArgs().front();
      indices.push_back(loop.getInductionVar());

      if (i == 0) {
        rewriter.replaceOp(op, loop);
      } else {
        affine::AffineYieldOp::create(rewriter, op->getLoc(),
                                      loop->getResults());
      }
      rewriter.setInsertionPointToStart(loop.getBody());
    }

    // Create the innermost body
    auto resultTypes =
        llvm::to_vector<6>(llvm::map_range(op->getResultTypes(), [](Type type) {
          if (auto tensorTy = dyn_cast<RankedTensorType>(type)) {
            return tensorTy.getElementType();
          }
          return type;
        }));

    // Generate a `tensor.extract` or `tensor.extract_slice` for each tensor
    // operand
    SmallVector<Value, 4> newOperands;
    for (OpOperand& opOperand : op->getOpOperands()) {
      Value operand = opOperand.get();
      auto rankedTensorTy = dyn_cast<RankedTensorType>(operand.getType());
      if (rankedTensorTy) {
        if (OpTrait::hasElementwiseMappableTraits(op)) {
          // We don't need to check the shape, as ElementwiseMappable
          // requires all tensor operands to have compatible shapes
          auto extractOp = tensor::ExtractOp::create(rewriter, operand.getLoc(),
                                                     operand, indices);
          newOperands.push_back(extractOp);
          continue;
        }

        auto interface = cast<ElementwiseByOperandOpInterface>(op);
        int operandIndex = opOperand.getOperandNumber();
        if (interface.operandIsMappable(operandIndex)) {
          auto mappedDims = interface.mappedDimensionsForOperand(operandIndex);
          if (mappedDims.size() == (size_t)rankedTensorTy.getRank()) {
            auto extractOp = tensor::ExtractOp::create(
                rewriter, operand.getLoc(), operand, indices);
            newOperands.push_back(extractOp);
            continue;
          }

          // Partial mapping: use extract_slice with rank reduction.
          // We assume mapped dimensions are a contiguous prefix.
          int64_t operandRank = rankedTensorTy.getRank();
          SmallVector<OpFoldResult> offsets, sizes, strides;
          for (int64_t d = 0; d < operandRank; ++d) {
            if (llvm::is_contained(mappedDims, d)) {
              offsets.push_back(indices[d]);
              sizes.push_back(rewriter.getIndexAttr(1));
            } else {
              offsets.push_back(rewriter.getIndexAttr(0));
              sizes.push_back(
                  rewriter.getIndexAttr(rankedTensorTy.getDimSize(d)));
            }
            strides.push_back(rewriter.getIndexAttr(1));
          }

          // The result type of the rank-reduced extract_slice
          SmallVector<int64_t> reducedShape;
          for (int64_t d = 0; d < operandRank; ++d) {
            if (!llvm::is_contained(mappedDims, d)) {
              reducedShape.push_back(rankedTensorTy.getDimSize(d));
            }
          }
          auto reducedType = RankedTensorType::get(
              reducedShape, rankedTensorTy.getElementType(),
              rankedTensorTy.getEncoding());

          auto extractSliceOp = tensor::ExtractSliceOp::create(
              rewriter, operand.getLoc(), reducedType, operand, offsets, sizes,
              strides);
          newOperands.push_back(extractSliceOp);
          continue;
        }
      }

      // We are left with either non-tensor operands (e.g., scalars) or tensor
      // operands which are not to be mapped over. These are replicated and
      // used as-is.
      newOperands.push_back(operand);
    }

    // "lowered" operation is the same operation, but over non-tensor
    // operands
    auto* scalarOp =
        rewriter.create(op->getLoc(), op->getName().getIdentifier(),
                        newOperands, resultTypes, op->getAttrs());

    // insert scalarOp into the tensor at right index
    Value inserted = tensor::InsertOp::create(
        rewriter, op->getLoc(), scalarOp->getResult(tensorResultIndex), target,
        indices);

    // replace lingalg.yield scalarOp with affine.yield insertedOp
    affine::AffineYieldOp::create(rewriter, op->getLoc(), inserted);
    return success();
  }
};
}  // namespace

struct ElementwiseToAffine
    : impl::ElementwiseToAffineBase<ElementwiseToAffine> {
  using ElementwiseToAffineBase::ElementwiseToAffineBase;

  void runOnOperation() override {
    MLIRContext* context = &getContext();

    // Walk to check for dynamic tensor shapes among ops that will be converted,
    // and fail if you find any.
    auto walkResult = getOperation()->walk([&](Operation* op) {
      bool convertAll = convertDialects.empty() && convertOps.empty();
      bool convertDialect = llvm::is_contained(
          convertDialects, op->getDialect()->getNamespace().str());
      bool convertOp =
          llvm::is_contained(convertOps, op->getName().getStringRef().str());

      if (convertAll || convertDialect || convertOp) {
        if (isSupported(op)) {
          SmallVector<Type> operandAndResultTypes;
          for (Type ty : op->getOperandTypes())
            operandAndResultTypes.push_back(ty);
          for (Type ty : op->getResultTypes())
            operandAndResultTypes.push_back(ty);

          for (Type ty : operandAndResultTypes) {
            ShapedType shapedTy = dyn_cast<ShapedType>(ty);
            if (shapedTy && !shapedTy.hasStaticShape()) {
              op->emitError()
                  << "op has operand or result with dynamic shape, "
                     "which is not supported by elementwise-to-affine";
              return WalkResult::interrupt();
            }
          }
        }
      }
      return WalkResult::advance();
    });

    if (walkResult.wasInterrupted()) {
      return signalPassFailure();
    }

    ConversionTarget target(*context);
    RewritePatternSet patterns(context);

    patterns.add<ConvertAnyElementwiseMappableOpOnRankedTensors>(context);
    target.markUnknownOpDynamicallyLegal([&](Operation* op) {
      bool convertAll = convertDialects.empty() && convertOps.empty();
      bool convertDialect = llvm::is_contained(
          convertDialects, op->getDialect()->getNamespace().str());
      bool convertOp =
          llvm::is_contained(convertOps, op->getName().getStringRef().str());

      if (convertAll || convertDialect || convertOp) return !isSupported(op);
      return true;
    });

    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns))))
      signalPassFailure();
  }
};

}  // namespace heir
}  // namespace mlir
