#include "lib/Transforms/LayoutOptimization/Patterns.h"

#include <optional>

#include "lib/Dialect/Secret/IR/SecretAttributes.h"
#include "lib/Dialect/Secret/IR/SecretDialect.h"
#include "lib/Dialect/Secret/IR/SecretOps.h"
#include "lib/Dialect/TensorExt/IR/TensorExtDialect.h"
#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
#include "lib/Kernel/KernelName.h"
#include "lib/Utils/AttributeUtils.h"
#include "lib/Utils/Layout/Utils.h"
#include "llvm/include/llvm/ADT/STLExtras.h"            // from @llvm-project
#include "llvm/include/llvm/ADT/SmallVector.h"          // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributes.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/Support/LLVM.h"             // from @llvm-project

namespace mlir {
namespace heir {

using tensor_ext::AssignLayoutOp;
using tensor_ext::ConvertLayoutOp;

namespace {
auto& kLayoutAttrName = tensor_ext::TensorExtDialect::kLayoutAttrName;
}

LogicalResult tryFoldLayoutConversionIntoPrevious(
    RewriterBase& rewriter, ConvertLayoutOp op,
    SmallVector<Operation*>& opsToErase) {
  if (auto priorConversion = op.getValue().getDefiningOp<ConvertLayoutOp>()) {
    // merge the two conversions into one
    auto newFrom = priorConversion.getFromLayoutAttr();
    auto newTo = op.getToLayoutAttr();
    auto newConversion = ConvertLayoutOp::create(
        rewriter, op.getLoc(), priorConversion.getValue(), newFrom, newTo,
        op.getDomainScheduleAttr());
    newConversion->setAttr(kLayoutAttrName, newTo);

    rewriter.replaceAllUsesWith(op, newConversion);
    opsToErase.push_back(op);
    if (priorConversion->use_empty()) {
      opsToErase.push_back(priorConversion);
    }
    return success();
  }

  if (auto priorAssignment = op.getValue().getDefiningOp<AssignLayoutOp>()) {
    if (isa<ArrayAttr>(priorAssignment.getLayout())) {
      // Don't fold into atomic or complex assignments to prevent performance
      // regressions.
      return failure();
    }
    // merge the conversion into the assignment return success();
    auto newLayout = op.getToLayoutAttr();
    auto domainSchedule = !op.getDomainSchedule().empty()
                              ? op.getDomainScheduleAttr()
                              : priorAssignment.getDomainScheduleAttr();
    auto newAssign = AssignLayoutOp::create(rewriter, op.getLoc(),
                                            priorAssignment.getValue(),
                                            newLayout, domainSchedule);
    newAssign->setAttr(kLayoutAttrName, newLayout);

    rewriter.replaceAllUsesWith(op, newAssign);
    opsToErase.push_back(op);
    if (priorAssignment->use_empty()) {
      opsToErase.push_back(priorAssignment);
    }
    return success();
  }

  return failure();
}

LogicalResult FoldLayoutConversions::matchAndRewrite(
    tensor_ext::ConvertLayoutOp op, PatternRewriter& rewriter) const {
  SmallVector<Operation*> opsToErase;
  LogicalResult result =
      tryFoldLayoutConversionIntoPrevious(rewriter, op, opsToErase);

  for (auto* op : opsToErase) {
    rewriter.eraseOp(op);
  }

  return result;
}

LogicalResult FoldMatvecInputConversionIntoPlaintext::matchAndRewrite(
    linalg::MatvecOp op, PatternRewriter& rewriter) const {
  auto conversion = op.getInputs()[1].getDefiningOp<ConvertLayoutOp>();
  auto matrixAssignment = op.getInputs()[0].getDefiningOp<AssignLayoutOp>();
  // Don't fold into atomic or complex (ArrayAttr) assignments to prevent
  // performance regressions.
  if (!conversion || !matrixAssignment ||
      isa<ArrayAttr>(matrixAssignment.getLayout())) {
    return failure();
  }

  auto fromLayout =
      dyn_cast<tensor_ext::LayoutAttr>(conversion.getFromLayout());
  auto toLayout = dyn_cast<tensor_ext::LayoutAttr>(conversion.getToLayout());
  auto matrixLayout =
      dyn_cast<tensor_ext::LayoutAttr>(matrixAssignment.getLayout());
  auto vectorType = dyn_cast<RankedTensorType>(conversion.getType());
  auto matrixType = dyn_cast<RankedTensorType>(matrixAssignment.getType());
  if (!fromLayout || !toLayout || !matrixLayout || !vectorType || !matrixType) {
    return failure();
  }

  auto kernel = op->getAttrOfType<secret::KernelAttr>(
      secret::SecretDialect::kKernelAttrName);
  if (!kernel || kernel.getName() != KernelName::MatvecDiagonal)
    return failure();

  const auto& toRelation = toLayout.getIntegerRelation();
  auto ciphertextSize = toRelation.getConstantBoundOnDimSize64(
      toRelation.getVarKindOffset(presburger::VarKind::Range) + 1);
  if (!ciphertextSize ||
      !isRelationRowMajor(vectorType, *ciphertextSize, toRelation) ||
      !isRelationSquatDiagonal(matrixType, *ciphertextSize,
                               matrixLayout.getIntegerRelation()) ||
      !isOneToOneSingleCiphertextPacking(fromLayout.getIntegerRelation())) {
    return failure();
  }

  auto newMatrixLayout = tensor_ext::LayoutAttr::getFromIntegerRelation(
      op.getContext(),
      foldVectorPermutationIntoMatrixLayout(fromLayout.getIntegerRelation(),
                                            matrixLayout.getIntegerRelation()));

  // Emit the repacked (plaintext) weights where the original assignment lived,
  // rather than inside the secret.generic that contains the matvec.
  rewriter.setInsertionPoint(matrixAssignment);
  auto newMatrix = AssignLayoutOp::create(
      rewriter, op.getLoc(), matrixAssignment.getValue(), newMatrixLayout,
      matrixAssignment.getDomainScheduleAttr());
  newMatrix->setAttr(kLayoutAttrName, newMatrixLayout);

  rewriter.modifyOpInPlace(op, [&] {
    op->setOperand(0, newMatrix);
    op->setOperand(1, conversion.getValue());
  });
  if (conversion->use_empty()) rewriter.eraseOp(conversion);
  if (matrixAssignment->use_empty()) rewriter.eraseOp(matrixAssignment);
  return success();
}

LogicalResult HoistArgLayouts::matchAndRewrite(
    func::FuncOp func, PatternRewriter& rewriter) const {
  bool changed = false;

  auto getFirstLayoutConversionOp =
      [&](OpOperand& use) -> std::optional<ConvertLayoutOp> {
    auto genericOp = dyn_cast<secret::GenericOp>(use.getOwner());
    if (!genericOp) {
      // A block argument may be returned directly, so no layout conversion is
      // performed.
      return std::nullopt;
    }
    auto innerArg = genericOp.getBody(0)->getArgument(use.getOperandNumber());
    ConvertLayoutOp convertLayoutOp = nullptr;
    for (auto user : innerArg.getUsers()) {
      if (user->getDialect()->getNamespace() == "debug") {
        continue;
      }
      // Any further layout conversions after operations that don't change
      // layout were already hoisted through in the layout-optimization pass.
      if (auto op = dyn_cast<ConvertLayoutOp>(user)) {
        if (convertLayoutOp) return std::nullopt;
        convertLayoutOp = op;
      } else {
        return std::nullopt;
      }
    }
    if (!convertLayoutOp) return std::nullopt;

    return convertLayoutOp;
  };

  for (auto& blockArg : func.getArguments()) {
    // Check that all uses of the block argument have the same layout
    // conversion. Otherwise, hoisting may not produce a benefit; it would
    // require duplicating the function argument or updating the layout
    // conversions of other uses.
    auto maybeLayoutOps =
        llvm::map_range(blockArg.getUses(), getFirstLayoutConversionOp);
    if (maybeLayoutOps.empty()) continue;
    auto layoutOpsVec = llvm::to_vector(maybeLayoutOps);
    auto maybeLayouts = llvm::map_range(
        layoutOpsVec, [](auto& maybeLayoutOp) -> std::optional<Attribute> {
          if (!maybeLayoutOp.has_value()) return std::nullopt;
          auto toLayout = maybeLayoutOp.value().getToLayoutAttr();
          if (!toLayout) return std::nullopt;
          return toLayout;
        });
    if (!llvm::all_equal(maybeLayouts)) continue;

    // If there was no layout conversion, there is nothing to hoist.
    auto maybeLayoutOp = *maybeLayoutOps.begin();
    if (!maybeLayoutOp.has_value()) continue;

    Attribute toLayout = maybeLayoutOp.value().getToLayoutAttr();

    rewriter.modifyOpInPlace(func, [&]() {
      // Update the function argument layout and the layout conversion's input
      // attribute
      func.setArgAttr(blockArg.getArgNumber(), kLayoutAttrName, toLayout);
      for (auto maybeLayoutOp : maybeLayoutOps) {
        auto innerValue = maybeLayoutOp.value().getValue();
        setAttributeAssociatedWith(innerValue, kLayoutAttrName, toLayout);
      }
    });
    for (auto maybeLayoutOp : maybeLayoutOps) {
      rewriter.replaceOp(maybeLayoutOp.value(),
                         {maybeLayoutOp.value().getValue()});
    }
    changed = true;
  }
  return changed ? success() : failure();
}

}  // namespace heir
}  // namespace mlir
