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

#include <cmath>
#include <cstdint>
#include <optional>
#include <string>

#include "lib/Analysis/LevelAnalysis/LevelAnalysis.h"
#include "lib/Analysis/MulDepthAnalysis/MulDepthAnalysis.h"
#include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h"
#include "lib/Dialect/Mgmt/IR/MgmtOps.h"
#include "llvm/include/llvm/ADT/STLExtras.h"               // from @llvm-project
#include "llvm/include/llvm/Support/Debug.h"               // from @llvm-project
#include "llvm/include/llvm/Support/DebugLog.h"            // from @llvm-project
#include "mlir/include/mlir/Analysis/DataFlowFramework.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/LoopUtils.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
#include "mlir/include/mlir/Dialect/SCF/Transforms/Transforms.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/SCF/Utils/Utils.h"  // from @llvm-project
#include "mlir/include/mlir/IR/Attributes.h"            // from @llvm-project
#include "mlir/include/mlir/IR/Block.h"                 // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributes.h"     // from @llvm-project
#include "mlir/include/mlir/IR/Matchers.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/ValueRange.h"            // from @llvm-project
#include "mlir/include/mlir/Interfaces/ControlFlowInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/Interfaces/LoopLikeInterface.h"  // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project

#define DEBUG_TYPE "halo-patterns"

namespace mlir {
namespace heir {

using affine::AffineForOp;
using affine::loopUnrollByFactor;

// Test if the two lists of values have mismatching secretness. Nb., a secret
// init with a non-secret iter arg is probably not a meaningful outcome of this
// function, but it doesn't hurt to support it.
static bool hasMismatch(ArrayRef<Value> inits, ArrayRef<Value> yieldedValues,
                        DataFlowSolver* solver) {
  for (int i = 0; i < inits.size(); ++i) {
    Value init = inits[i];
    Value yield = yieldedValues[i];

    auto* initLattice = solver->lookupState<SecretnessLattice>(init);
    if (!initLattice || !initLattice->getValue().isInitialized()) {
      continue;
    }

    auto* yieldLattice = solver->lookupState<SecretnessLattice>(yield);
    if (!yieldLattice || !yieldLattice->getValue().isInitialized()) {
      continue;
    }

    if (initLattice->getValue().getSecretness() !=
        yieldLattice->getValue().getSecretness()) {
      return true;
    }
  }

  return false;
}

static bool isInsideLoopNesting(Operation* op, Block* block) {
  if (op->getBlock() == block) {
    return false;
  }
  Operation* parent = op->getParentOp();
  while (parent && parent->getBlock() != block) {
    if (isa<LoopLikeOpInterface>(parent)) {
      return true;
    }
    parent = parent->getParentOp();
  }
  if (parent && isa<LoopLikeOpInterface>(parent)) {
    return true;
  }
  return false;
}

LogicalResult PeelPlaintextAffineForInit::matchAndRewrite(
    AffineForOp forOp, PatternRewriter& rewriter) const {
  // Determine if the inits are public while the yielded values are secret
  SmallVector<Value> inits(forOp.getInits());
  SmallVector<Value> yieldedValues(forOp.getYieldedValues());
  bool mismatch = hasMismatch(inits, yieldedValues, solver);

  if (!mismatch) {
    return rewriter.notifyMatchFailure(
        forOp, "No ciphertext/plaintext mismatch detected");
  }

  LLVM_DEBUG(
      llvm::dbgs() << "Found mismatch between init types and iter arg types\n");

  if (!forOp.hasConstantBounds())
    return rewriter.notifyMatchFailure(
        forOp, "Only affine.for ops with constant bounds are supported");

  int64_t lb = forOp.getConstantLowerBound();
  int64_t ub = forOp.getConstantUpperBound();
  int64_t step = forOp.getStep().getSExtValue();
  int64_t splitBound = lb + step;

  if (ceil(float(ub - lb) / step) <= 1) {
    return rewriter.notifyMatchFailure(
        forOp, "Peeling is not needed if there is one or less iteration.");
  }

  // Split the loop into two loops, the first of which has just one iteration,
  // and then unroll the single-iteration loop.
  AffineForOp firstIteration =
      cast<AffineForOp>(rewriter.clone(*forOp.getOperation()));
  rewriter.modifyOpInPlace(firstIteration, [&]() {
    firstIteration.setConstantUpperBound(splitBound);
  });

  rewriter.modifyOpInPlace(forOp, [&]() {
    forOp.getInitsMutable().assign(firstIteration->getResults());
    forOp.setConstantLowerBound(splitBound);
  });

  if (failed(loopUnrollFull(firstIteration))) {
    LLVM_DEBUG(llvm::dbgs()
               << "Failed to unroll single-iteration affine.for!\n");
  }
  return success();
}

LogicalResult PeelPlaintextScfForInit::matchAndRewrite(
    scf::ForOp forOp, PatternRewriter& rewriter) const {
  // Determine if the inits are public while the yielded values are secret
  SmallVector<Value> inits(forOp.getInits());
  SmallVector<Value> yieldedValues(forOp.getYieldedValues());
  bool mismatch = hasMismatch(inits, yieldedValues, solver);

  if (!mismatch) {
    return rewriter.notifyMatchFailure(
        forOp, "No ciphertext/plaintext mismatch detected");
  }

  LLVM_DEBUG(
      llvm::dbgs() << "Found mismatch between init types and iter arg types\n");

  scf::ForOp firstIteration;
  LogicalResult result =
      peelForLoopFirstIteration(rewriter, forOp, firstIteration);
  if (failed(result)) {
    return failure();
  }

  if (failed(loopUnrollFull(firstIteration))) {
    LLVM_DEBUG(llvm::dbgs() << "Failed to unroll single-iteration scf.for!\n");
  }

  return success();
}

// Helper for rejecting a loop unroll pattern. Check:
//
// - all secret iter args are immediately bootstrapped and have no other uses
// - all secret iter args level_reduce_min just before being yielded
//
// If successful, the SSA values corresponding to secret iter args are returned.
template <typename ForLoop>
FailureOr<SmallVector<Value>> isLoopStructuredForHaloUnroll(
    ForLoop forOp, DataFlowSolver* solver) {
  LLVM_DEBUG(
      llvm::dbgs() << "Checking loop preconditions for Halo loop unroll\n");
  SmallVector<Value> secretIterArgs;

  for (Value iterArg : forOp.getRegionIterArgs()) {
    if (isSecret(iterArg, solver)) {
      if (!iterArg.hasOneUse()) {
        LDBG(2) << "Iter arg " << iterArg << " has " << iterArg.getNumUses()
                << " uses, expected 1";
        return failure();
      }
      auto* user = *iterArg.getUsers().begin();
      if (!isa<mgmt::BootstrapOp>(user)) {
        LDBG(2) << "Iter arg " << iterArg
                << " single use is not a bootstrap op, instead "
                << user->getName();
        return failure();
      }

      Value yieldedValue =
          forOp.getTiedLoopYieldedValue(cast<BlockArgument>(iterArg))->get();
      if (!isa<mgmt::LevelReduceMinOp>(yieldedValue.getDefiningOp())) {
        LDBG(2) << "Yielded value " << yieldedValue
                << " is not the result of a level_reduce_min op, instead "
                << yieldedValue.getDefiningOp()->getName();
        return failure();
      }

      secretIterArgs.push_back(iterArg);
    } else {
      LDBG(2) << "Iter arg " << iterArg << " is not secret";
    }
  }

  if (secretIterArgs.empty()) {
    LDBG(2) << "No secret iter args found";
    return failure();
  }

  return secretIterArgs;
}

template <typename ForOp>
LogicalResult doPartialUnroll(ForOp forOp, PatternRewriter& rewriter,
                              int forceMaxLevel, DataFlowSolver* solver) {
  LDBG(2) << "In doPartialUnroll";
  FailureOr<SmallVector<Value>> secretIterArgsResult =
      isLoopStructuredForHaloUnroll(forOp, solver);
  if (failed(secretIterArgsResult)) {
    LDBG(2) << "Loop preconditions not met";
    return rewriter.notifyMatchFailure(forOp, "Loop preconditions not met");
  }
  LDBG(2) << "Loop is structured for Halo-style unroll";
  SmallVector<Value> secretIterArgs = secretIterArgsResult.value();

  std::optional<APInt> maybeTripCount = forOp.getStaticTripCount();
  if (!maybeTripCount.has_value()) {
    LDBG(2) << "Could not infer trip count";
    return failure();
  }
  int64_t tripCount = maybeTripCount->getSExtValue();
  LDBG(2) << "Inferred trip count=" << tripCount;

  // Now we need to compute how many loop iterations we can unroll.
  //
  // To start, assume we have just one iter arg. It is assumed to be the result
  // of a mgmt.level_reduce_min op just before being yielded. The quantity we
  // need to compute is the difference between the level just after the
  // bootstrap op (the "effective level" of bootstrap) and the level just before
  // the level_reduce_min op.
  //
  //  %2 = mgmt.level_reduce_min %1 : i32
  //  %3 = affine.for %arg1 = 1 to 12 iter_args(%arg2 = %2) -> (i32) {
  //     %4 = mgmt.bootstrap %arg2 : i32
  //     LEVEL_START[%arg2] = level(%4)
  //
  //     ...
  //
  //     LEVEL_END[%arg2] = level(%5)
  //     %6 = mgmt.level_reduce_min %5 : i32
  //     affine.yield %6 : i32
  //  }
  //
  // If LEVEL_END > 0, then we can unroll if and only if LEVEL_START / T > 1,
  // and (int) LEVEL_START / T is the factor we can unroll by.
  //
  // If we have many iter args, the allowable loop unroll factor is the minimum
  // among all iter args.
  //
  // TODO(#2556): Use smarter unroll strategies for multiple iter_args.

  SmallVector<int> unrollFactors;
  for (Value iterArg : secretIterArgs) {
    auto bootstrapOp = cast<mgmt::BootstrapOp>(*iterArg.getUsers().begin());
    auto* levelStart =
        solver->lookupState<LevelLattice>(bootstrapOp.getResult());

    auto yieldedValue =
        forOp.getTiedLoopYieldedValue(cast<BlockArgument>(iterArg))->get();
    mgmt::LevelReduceMinOp levelReduceOp =
        cast<mgmt::LevelReduceMinOp>(yieldedValue.getDefiningOp());
    auto* levelEnd =
        solver->lookupState<LevelLattice>(levelReduceOp.getInput());

    if (!levelStart || !levelEnd || !levelStart->getValue().isInt() ||
        !levelEnd->getValue().isInt()) {
      LDBG(2)
          << "Start and end levels were not inferable to be concrete integers";
      return rewriter.notifyMatchFailure(
          forOp,
          "Start and end levels were not inferable to be concrete integers");
    }

    int levelEndVal = levelEnd->getValue().getInt();
    int levelStartVal = levelStart->getValue().getInt();
    LDBG(2) << "levelEndVal=" << levelEndVal
            << ", levelStartVal=" << levelStartVal;

    LDBG(2) << "Getting levelAfterBootstrap";
    int levelAfterBootstrap = 0;

    if (forceMaxLevel > 0) {
      levelAfterBootstrap = forceMaxLevel;
      LDBG(2) << "Using forced max level of " << forceMaxLevel;
    } else {
      levelAfterBootstrap =
          getMaxLevel(forOp->template getParentOfType<func::FuncOp>(), solver);
    }
    int levelsUsedInLoop = levelEndVal - levelStartVal;

    LDBG(2) << "levelsUsedInLoop=" << levelsUsedInLoop << ", "
            << "levelAfterBootstrap=" << levelAfterBootstrap;

    if (levelsUsedInLoop == 0) {
      LDBG(2) << "Loop uses zero levels. iter_arg=" << iterArg
              << "; levelStartVal=" << levelStartVal
              << "; levelEndVal=" << levelEndVal
              << "; levelAfterBootstrap=" << levelAfterBootstrap
              << "; op was: " << forOp;
      return failure();
    }
    LDBG(2) << "Levels used in loop = " << levelsUsedInLoop;

    int unrollFactor = (levelAfterBootstrap - levelStartVal) / levelsUsedInLoop;
    LDBG(2) << "Found unroll factor " << unrollFactor
            << " for iter_arg=" << iterArg
            << "; levelStartVal=" << levelStartVal
            << "; levelEndVal=" << levelEndVal
            << "; levelAfterBootstrap=" << levelAfterBootstrap;
    if (unrollFactor > 1) {
      unrollFactors.push_back(unrollFactor);
    }
  }

  if (unrollFactors.empty()) {
    LDBG(2) << "No unroll factors found";
    return rewriter.notifyMatchFailure(forOp,
                                       "No unroll factors could be found.");
  }

  int chosenUnrollFactor = *llvm::min_element(unrollFactors);
  chosenUnrollFactor =
      tripCount < chosenUnrollFactor ? tripCount : chosenUnrollFactor;
  if (chosenUnrollFactor > 1) {
    Block* block = forOp->getBlock();
    // The function_ref<void(unsigned, Operation *, OpBuilder)> annotateFn that
    // we pass to the loop unroll step ensures that we can tell which bootstrap
    // ops and level_reduce_min ops are safe to remove post-unroll.
    LDBG(2) << "Applying Halo-style unroll by a factor of "
            << chosenUnrollFactor;

    // First mark the special loop invariance ops with an attribute, so that any
    // other operations of the same type in the loop body are not accidentally
    // deleted by cleanup. The name of the attribute is arbitrary as it will be
    // cleaned up before this pattern is complete.
    std::string specialOpKey = "halo.invariance";
    for (Value iterArg : forOp.getRegionIterArgs()) {
      if (isSecret(iterArg, solver)) {
        auto bootstrapOp = cast<mgmt::BootstrapOp>(*iterArg.getUsers().begin());
        LDBG(2) << "Marking op " << bootstrapOp
                << " with halo.invariance so it can be analyzed for removal";
        rewriter.modifyOpInPlace(bootstrapOp, [&]() {
          bootstrapOp->setAttr(specialOpKey, rewriter.getUnitAttr());
        });

        Value yieldedValue =
            forOp.getTiedLoopYieldedValue(cast<BlockArgument>(iterArg))->get();
        auto levelReduceOp =
            cast<mgmt::LevelReduceMinOp>(yieldedValue.getDefiningOp());
        LDBG(2) << "Marking op " << levelReduceOp
                << " with halo.invariance so it can be analyzed for removal";
        rewriter.modifyOpInPlace(levelReduceOp, [&]() {
          levelReduceOp->setAttr(specialOpKey, rewriter.getUnitAttr());
        });
      }
    }

    LDBG(2) << "Starting loop unroll";
    if (failed(loopUnrollByFactor(
            forOp, chosenUnrollFactor,
            [&](unsigned index, Operation* clonedOp, OpBuilder builder) {
              if (!clonedOp->hasAttr(specialOpKey)) return;

              // Internally, loopUnrollByFactor starts the passed `index` at 1
              // and goes up to unrollFactor (exclusive). This index 1
              // corresponds to the second loop iteration, but this callable is
              // never called for index 0 because the original loop body
              // is left intact and the additional unrolled iterations are
              // inserted after it. However, this is an implementation detail
              // that we try not to rely on.
              if (index > 0 && isa<mgmt::BootstrapOp>(clonedOp)) {
                LDBG(2) << "Marking op " << *clonedOp
                        << " with halo.remove so it can be removed";
                // We cannot remove the op in the middle of the loop unrolling
                // process, so instead mark it for later removal.
                rewriter.modifyOpInPlace(clonedOp, [&]() {
                  clonedOp->setAttr("halo.remove", builder.getUnitAttr());
                });
              } else if (index < chosenUnrollFactor - 1 &&
                         isa<mgmt::LevelReduceMinOp>(clonedOp)) {
                LDBG(2) << "Marking op " << *clonedOp
                        << " with halo.remove so it can be removed";
                rewriter.modifyOpInPlace(clonedOp, [&]() {
                  clonedOp->setAttr("halo.remove", builder.getUnitAttr());
                });
              }
              rewriter.modifyOpInPlace(
                  clonedOp, [&]() { clonedOp->removeAttr(specialOpKey); });
            })))
      return failure();

    // Clean up remainder iterations.
    SmallVector<mgmt::LevelReduceMinOp> opsToRemove;
    block->walk([&](mgmt::LevelReduceMinOp reduceOp) {
      if (reduceOp->hasAttr("halo.invariance") &&
          !isInsideLoopNesting(reduceOp, block)) {
        opsToRemove.push_back(reduceOp);
      }
    });
    for (auto reduceOp : opsToRemove) {
      LDBG(2) << "Removing leftover level_reduce_min: " << reduceOp;
      rewriter.replaceOp(reduceOp, reduceOp.getOperand());
    }
  }

  return success();
}

LogicalResult PartialUnrollForLevelConsumptionAffineFor::matchAndRewrite(
    affine::AffineForOp forOp, PatternRewriter& rewriter) const {
  return doPartialUnroll(forOp, rewriter, forceMaxLevel, solver);
}

LogicalResult PartialUnrollForLevelConsumptionSCFFor::matchAndRewrite(
    scf::ForOp forOp, PatternRewriter& rewriter) const {
  return doPartialUnroll(forOp, rewriter, forceMaxLevel, solver);
}

LogicalResult DeleteAnnotatedOps::matchAndRewrite(
    Operation* op, PatternRewriter& rewriter) const {
  if (!op->hasAttrOfType<UnitAttr>("halo.remove")) {
    return failure();
  }

  if (op->getNumResults() != 1 || op->getNumOperands() != 1) {
    return failure();
  }

  rewriter.replaceOp(op, op->getOperand(0));
  return success();
}

OpOperand* getOpOperandForwardingToOpResult(RegionBranchOpInterface op,
                                            Value predecessor,
                                            OpResult result) {
  for (auto user : predecessor.getUsers()) {
    if (auto terminatorOp = dyn_cast<RegionBranchTerminatorOpInterface>(user)) {
      OperandRange succOperands =
          terminatorOp.getSuccessorOperands(RegionSuccessor(op.getOperation()));
      unsigned operandIndex =
          succOperands.getBeginOperandIndex() + result.getResultNumber();
      return &terminatorOp->getOpOperand(operandIndex);
    }
  }
  return nullptr;
}

LogicalResult RegionBranchOpLevelInvariancePattern::matchAndRewrite(
    RegionBranchOpInterface op, PatternRewriter& rewriter) const {
  bool changed = false;

  // For each result of the op, check if the levels of all yielding branches
  // are invariant.
  for (int i = 0; i < op->getNumResults(); ++i) {
    Value result = op->getResult(i);
    if (!isSecret(result, solver)) continue;

    // The mapping produces for each Value a SmallVector of OpOperand* that can
    // flow into this specific result.
    mlir::RegionBranchInverseSuccessorMapping inverseMapping;
    op.getSuccessorInputOperandMapping(inverseMapping);
    llvm::SmallVector<mlir::OpOperand*> yieldingOperands;
    yieldingOperands = inverseMapping.lookup(result);

    if (yieldingOperands.size() < 2) continue;

    // Collect levels for all yielding operands.
    SmallVector<int64_t> levels;
    for (OpOperand* operand : yieldingOperands) {
      auto* levelLattice = solver->lookupState<LevelLattice>(operand->get());
      if (!levelLattice || !levelLattice->getValue().isInt()) {
        // If level is not an int, we cannot easily reconcile it.
        continue;
      }
      levels.push_back(levelLattice->getValue().getInt());
    }

    // Reconcile levels if we have all of them.
    if (levels.size() == yieldingOperands.size()) {
      // The context LevelAnalysis is only a forward analysis, which does
      // levels from 0 ... L, since the absolute max level is not yet known.
      int64_t maxLevel = *llvm::max_element(levels);
      for (int j = 0; j < yieldingOperands.size(); ++j) {
        if (levels[j] < maxLevel) {
          int64_t diff = maxLevel - levels[j];
          rewriter.setInsertionPoint(yieldingOperands[j]->getOwner());
          auto reduceOp = mgmt::LevelReduceOp::create(
              rewriter, op.getLoc(), yieldingOperands[j]->get(), diff);
          rewriter.modifyOpInPlace(yieldingOperands[j]->getOwner(), [&]() {
            yieldingOperands[j]->set(reduceOp.getResult());
          });
          changed = true;
        }
      }
    }
  }

  return changed ? success() : failure();
}

LogicalResult RegionBranchOpScaleInvariancePattern::matchAndRewrite(
    RegionBranchOpInterface op, PatternRewriter& rewriter) const {
  bool changed = false;

  mlir::RegionBranchInverseSuccessorMapping inverseMapping;
  op.getSuccessorInputOperandMapping(inverseMapping);
  for (int i = 0; i < op->getNumResults(); ++i) {
    Value result = op->getResult(i);
    if (!isSecret(result, solver)) continue;

    llvm::SmallVector<mlir::OpOperand*> yieldingOperands;
    yieldingOperands = inverseMapping.lookup(result);

    if (yieldingOperands.size() < 2) continue;

    SmallVector<int64_t> mulDepths;
    for (OpOperand* operand : yieldingOperands) {
      auto* mulDepthLattice =
          solver->lookupState<MulDepthLattice>(operand->get());
      if (!mulDepthLattice || !mulDepthLattice->getValue().isInt()) {
        continue;
      }
      mulDepths.push_back(mulDepthLattice->getValue().getMulDepth());
    }

    // Reconcile mul depths if we have all of them.
    if (mulDepths.size() == yieldingOperands.size()) {
      int64_t maxMulDepth = *llvm::max_element(mulDepths);
      for (int j = 0; j < yieldingOperands.size(); ++j) {
        if (mulDepths[j] < maxMulDepth) {
          rewriter.setInsertionPoint(yieldingOperands[j]->getOwner());
          auto adjustScaleOp = mgmt::AdjustScaleOp::create(
              rewriter, op.getLoc(), yieldingOperands[j]->get(),
              rewriter.getI64IntegerAttr((*idCounter)++));
          rewriter.modifyOpInPlace(yieldingOperands[j]->getOwner(), [&]() {
            yieldingOperands[j]->set(adjustScaleOp.getResult());
          });
          changed = true;
        }
      }
    }
  }

  return changed ? success() : failure();
}

LogicalResult doPartialUnroll(affine::AffineForOp forOp,
                              PatternRewriter& rewriter, int forceMaxLevel,
                              DataFlowSolver* solver) {
  return doPartialUnroll<affine::AffineForOp>(forOp, rewriter, forceMaxLevel,
                                              solver);
}

LogicalResult doPartialUnroll(scf::ForOp forOp, PatternRewriter& rewriter,
                              int forceMaxLevel, DataFlowSolver* solver) {
  return doPartialUnroll<scf::ForOp>(forOp, rewriter, forceMaxLevel, solver);
}

}  // namespace heir
}  // namespace mlir
