#ifndef LIB_KERNEL_KERNELIMPLEMENTATION_H_
#define LIB_KERNEL_KERNELIMPLEMENTATION_H_

#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <vector>

#include "lib/Kernel/AbstractValue.h"
#include "lib/Kernel/ArithmeticDag.h"
#include "lib/Kernel/KernelName.h"
#include "lib/Utils/APIntUtils.h"
#include "lib/Utils/MathUtils.h"
#include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace heir {
namespace kernel {

// A function that generalizes the reduction operation in all kernels in this
// file. E.g., whether to use `add` or `mul`
template <typename T>
using DagReducer = std::function<std::shared_ptr<ArithmeticDagNode<T>>(
    std::shared_ptr<ArithmeticDagNode<T>>,
    std::shared_ptr<ArithmeticDagNode<T>>)>;

template <typename T>
using DagExtractor = std::function<std::shared_ptr<ArithmeticDagNode<T>>(
    std::shared_ptr<ArithmeticDagNode<T>>, int64_t)>;

// Dynamic extraction: takes tensor and DAG node representing runtime index
template <typename T>
using DagExtractorDynamic = std::function<std::shared_ptr<ArithmeticDagNode<T>>(
    std::shared_ptr<ArithmeticDagNode<T>>,
    std::shared_ptr<ArithmeticDagNode<T>>)>;

// Returns an arithmetic DAG that implements a matvec kernel. Ensure this is
// only generated for T a subclass of AbstractValue.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementMatvec(KernelName kernelName, const T& matrix, const T& vector) {
  using NodeTy = ArithmeticDagNode<T>;
  assert(kernelName == KernelName::MatvecDiagonal);
  auto matrixDag = NodeTy::leaf(matrix);
  auto vectorDag = NodeTy::leaf(vector);

  int numRows = matrix.getShape()[0];
  assert(numRows > 0);

  auto firstTerm = NodeTy::mul(NodeTy::leftRotate(vectorDag, 0),
                               NodeTy::extract(matrixDag, 0));

  auto accumulatedSum = firstTerm;
  for (int i = 1; i < numRows; ++i) {
    auto term = NodeTy::mul(NodeTy::leftRotate(vectorDag, i),
                            NodeTy::extract(matrixDag, i));
    accumulatedSum = NodeTy::add(accumulatedSum, term);
  }
  return accumulatedSum;
}

// Returns an arithmetic DAG that implements a logarithmic rotate-and-reduce
// accumulation of an input ciphertext.
//
// This is a special case of `tensor_ext.rotate_and_reduce`
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementRotateAndReduceAccumulation(
    std::shared_ptr<ArithmeticDagNode<T>> vectorDag, int64_t period,
    int64_t steps, DagReducer<T> reduceFunc) {
  using NodeTy = ArithmeticDagNode<T>;
  for (int64_t shiftSize = steps / 2; shiftSize > 0; shiftSize /= 2) {
    auto rotated = NodeTy::leftRotate(vectorDag, shiftSize * period);
    auto reduced = reduceFunc(vectorDag, rotated);
    vectorDag = reduced;
  }
  return vectorDag;
}

template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementRotateAndReduceAccumulation(const T& vector, int64_t period,
                                     int64_t steps, DagReducer<T> reduceFunc) {
  using NodeTy = ArithmeticDagNode<T>;
  return implementRotateAndReduceAccumulation<T>(NodeTy::leaf(vector), period,
                                                 steps, reduceFunc);
}

// Rolled version of implementRotateAndReduceAccumulation.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementRotateAndReduceAccumulationRolled(
    std::shared_ptr<ArithmeticDagNode<T>> vectorDag, int64_t period,
    int64_t steps, DagReducer<T> reduceFunc, const DagType& baseType) {
  using NodeTy = ArithmeticDagNode<T>;
  using NodePtr = std::shared_ptr<NodeTy>;

  int64_t numIterations = static_cast<int64_t>(std::log2(steps));
  if (numIterations <= 0) return vectorDag;

  auto initialShift = NodeTy::constantScalar(steps / 2, DagType::index());
  auto loopNode = NodeTy::loop(
      {vectorDag, initialShift}, {baseType, DagType::index()}, /*lower=*/0,
      /*upper=*/numIterations,
      /*step=*/1, [&](NodePtr i, const std::vector<NodePtr>& iterArgs) {
        auto currentVector = iterArgs[0];
        auto currentShift = iterArgs[1];

        auto rotated = NodeTy::leftRotate(
            currentVector,
            NodeTy::mul(currentShift,
                        NodeTy::constantScalar(period, DagType::index())));
        auto reduced = reduceFunc(currentVector, rotated);

        auto nextShift = NodeTy::floorDiv(currentShift, 2);
        return NodeTy::yield({reduced, nextShift});
      });

  return NodeTy::resultAt(loopNode, 0);
}

// Rolled version of implementRotateAndReduceAccumulation.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementRotateAndReduceAccumulationRolled(const T& vector, int64_t period,
                                           int64_t steps,
                                           DagReducer<T> reduceFunc,
                                           const DagType& baseType) {
  using NodeTy = ArithmeticDagNode<T>;
  return implementRotateAndReduceAccumulationRolled<T>(
      NodeTy::leaf(vector), period, steps, reduceFunc, baseType);
}

// A function that generalizes the choice of rotation for the "baby stepped
// operand" of a baby-step giant-step algorithm. This is required because
// the rotation used in Halevi-Shoup matvec differs from that of bicyclic
// matmul.
using DerivedRotationIndexFn = std::function<int64_t(
    // giant step size
    int64_t,
    // current giant step index
    int64_t,
    // current baby step index
    int64_t,
    // period
    int64_t)>;

inline int64_t defaultDerivedRotationIndexFn(int64_t giantStepSize,
                                             int64_t giantStepIndex,
                                             int64_t babyStepIndex,
                                             int64_t period) {
  return -giantStepSize * giantStepIndex * period;
}

// Dynamic version: builds DAG expression for rotation amount from DAG node
// indices
template <typename T>
using DagDerivedRotationIndexFn =
    std::function<std::shared_ptr<ArithmeticDagNode<T>>(
        // giant step size (constant)
        int64_t,
        // current giant step index (DAG node)
        std::shared_ptr<ArithmeticDagNode<T>>,
        // current baby step index (DAG node)
        std::shared_ptr<ArithmeticDagNode<T>>,
        // period (constant)
        int64_t)>;

template <typename T>
std::shared_ptr<ArithmeticDagNode<T>> defaultDagDerivedRotationIndexFn(
    int64_t giantStepSize, std::shared_ptr<ArithmeticDagNode<T>> giantStepIndex,
    std::shared_ptr<ArithmeticDagNode<T>> babyStepIndex, int64_t period) {
  using NodeTy = ArithmeticDagNode<T>;
  // Build: -(giantStepSize * giantStepIndex * period)
  auto gsSize = NodeTy::constantScalar(giantStepSize, DagType::index());
  auto periodNode = NodeTy::constantScalar(period, DagType::index());
  auto negOne = NodeTy::constantScalar(-1, DagType::index());

  auto temp = NodeTy::mul(giantStepIndex, gsSize);
  temp = NodeTy::mul(temp, periodNode);
  return NodeTy::mul(temp, negOne);
}

// Returns an arithmetic DAG that implements a baby-step-giant-step
// rotate-and-reduce accumulation between an input ciphertext
// (giantSteppedOperand) and an abstraction over the other argument
// (babySteppedOperand). In particular, the babySteppedOperand may be a list of
// plaintexts like in Halevi-Shoup matvec, or a single ciphertext like in
// bicyclic matmul, and this abstracts over both by taking in an extraction
// callback.
//
// This is a special case of `tensor_ext.rotate_and_reduce`, but with the added
// abstractions it also supports situations not currently expressible by
// `tensor_ext.rotate_and_reduce`.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementBabyStepGiantStep(
    const T& giantSteppedOperand, const T& babySteppedOperand, int64_t period,
    int64_t steps, DagType dagType, DagExtractor<T> extractFunc,
    const std::map<int, bool>& zeroDiagonals = {},
    const DerivedRotationIndexFn& derivedRotationIndexFn =
        defaultDerivedRotationIndexFn) {
  using NodeTy = ArithmeticDagNode<T>;
  auto giantSteppedDag = NodeTy::leaf(giantSteppedOperand);
  auto babySteppedDag = NodeTy::leaf(babySteppedOperand);

  // Use a value of sqrt(n) as the baby step / giant step size.
  int64_t numBabySteps = static_cast<int64_t>(std::ceil(std::sqrt(steps)));
  int64_t giantStepSize = numBabySteps;
  // numGiantSteps = ceil(steps / numBabySteps)
  int64_t numGiantSteps = (steps + numBabySteps - 1) / numBabySteps;

  // Compute sqrt(n) ciphertext rotations of the input as baby-steps.
  SmallVector<std::shared_ptr<NodeTy>> babyStepVals;
  babyStepVals.push_back(giantSteppedDag);  // rot by zero
  for (int64_t i = 1; i < numBabySteps; ++i) {
    babyStepVals.push_back(NodeTy::leftRotate(giantSteppedDag, period * i));
  }

  // Compute the inner baby step sums.
  std::shared_ptr<NodeTy> result = nullptr;
  for (int64_t j = 0; j < numGiantSteps; ++j) {
    std::shared_ptr<NodeTy> innerSum = nullptr;
    for (int64_t i = 0; i < numBabySteps; ++i) {
      if (j * giantStepSize + i >= steps) {
        break;
      }
      int64_t innerRotAmount =
          derivedRotationIndexFn(giantStepSize, j, i, period);
      size_t extractionIndex = i + j * giantStepSize;

      // Skip the multiplication if the extraction index is zero.
      if (zeroDiagonals.contains(extractionIndex)) {
        continue;
      }

      auto plaintext = extractFunc(babySteppedDag, extractionIndex);
      auto rotatedPlaintext = NodeTy::leftRotate(plaintext, innerRotAmount);
      auto multiplied = NodeTy::mul(rotatedPlaintext, babyStepVals[i]);
      innerSum =
          innerSum == nullptr ? multiplied : NodeTy::add(innerSum, multiplied);
    }

    // The innerSum may be nullptr if all the multiplications were skipped.
    auto rotatedSum =
        innerSum == nullptr
            ? nullptr
            : NodeTy::leftRotate(innerSum, period * j * giantStepSize);
    if (result == nullptr) {
      result = rotatedSum;
    } else {
      result = rotatedSum == nullptr ? result : NodeTy::add(result, rotatedSum);
    }
  }

  return result == nullptr ? NodeTy::splat(0, dagType) : result;
}

// Default dynamic extractor: simple extraction at runtime index
template <typename T>
std::shared_ptr<ArithmeticDagNode<T>> defaultDagExtractor(
    std::shared_ptr<ArithmeticDagNode<T>> tensor,
    std::shared_ptr<ArithmeticDagNode<T>> index) {
  return ArithmeticDagNode<T>::extract(tensor, index);
}

// Rolled version of Baby-Step-Giant-Step algorithm.
//
// TODO(#2704): support skipping over zero diagonals
//
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementBabyStepGiantStepRolled(
    const T& giantSteppedOperand, const T& babySteppedOperand, int64_t period,
    int64_t steps, const DagType baseType,
    DagExtractorDynamic<T> extractFunc = defaultDagExtractor<T>,
    const std::map<int, bool>& zeroDiagonals = {},
    const DagDerivedRotationIndexFn<T>& dagRotationFn =
        defaultDagDerivedRotationIndexFn<T>) {
  using NodeTy = ArithmeticDagNode<T>;
  using NodePtr = std::shared_ptr<NodeTy>;

  auto giantSteppedDag = NodeTy::leaf(giantSteppedOperand);
  auto babySteppedDag = NodeTy::leaf(babySteppedOperand);

  int64_t numBabySteps = static_cast<int64_t>(std::ceil(std::sqrt(steps)));
  int64_t giantStepSize = numBabySteps;
  int64_t numGiantSteps = (steps + numBabySteps - 1) / numBabySteps;

  // Create an empty tensor of shape {numBabySteps, numSlots} and emit a loop to
  // compute baby-step rotations.
  std::vector<int64_t> packedShape = {numBabySteps};
  if (!baseType.getShape().empty()) {
    packedShape.push_back(baseType.getShape().back());
  }
  auto emptyTensor = NodeTy::empty(packedShape, baseType);
  auto precomputeLoop = NodeTy::loop(
      {emptyTensor}, {baseType}, /*lower=*/0, /*upper=*/numBabySteps,
      /*step=*/1, [&](NodePtr i, const std::vector<NodePtr>& iterArgs) {
        auto currentTensor = iterArgs[0];
        auto rotAmount =
            NodeTy::mul(i, NodeTy::constantScalar(period, DagType::index()));
        auto rotated = NodeTy::leftRotate(giantSteppedDag, rotAmount);
        auto updatedTensor = NodeTy::insert(rotated, currentTensor, i);
        return NodeTy::yield({updatedTensor});
      });
  NodePtr packedBabySteps = NodeTy::resultAt(precomputeLoop, 0);

  // Initialize outer sum to zero and pass packedBabySteps as loop invariant
  std::vector<NodePtr> outerInits = {NodeTy::splat(0, baseType),
                                     packedBabySteps};
  std::vector<DagType> outerTypes = {baseType, baseType};

  NodePtr isNotZeroDiagTensor = nullptr;
  if (!zeroDiagonals.empty()) {
    std::vector<double> isNotZeroDiagVec(steps, 1.0);
    for (const auto& [idx, isZero] : zeroDiagonals) {
      if (isZero && idx >= 0 && idx < steps) {
        isNotZeroDiagVec[idx] = 0.0;
      }
    }
    isNotZeroDiagTensor = NodeTy::constantTensor(
        isNotZeroDiagVec, DagType::floatTensor(64, {steps}));
  }

  // Outer loop over giant steps (j = 0 to numGiantSteps)
  auto outerLoop = NodeTy::loop(
      outerInits, outerTypes, /*lower=*/0, /*upper=*/numGiantSteps, /*step=*/1,
      [&](NodePtr j, const std::vector<NodePtr>& outerIterArgs) {
        auto outerSum = outerIterArgs[0];
        // Inner loop over baby steps (i = 0 to numBabySteps)
        auto innerLoop = NodeTy::loop(
            {NodeTy::splat(0, baseType)}, {baseType}, /*lower=*/0,
            /*upper=*/numBabySteps, /*step=*/1,
            [&](NodePtr i, const std::vector<NodePtr>& innerIterArgs) {
              auto innerSum = innerIterArgs[0];
              auto gsSize =
                  NodeTy::constantScalar(giantStepSize, DagType::index());
              auto jOffset = NodeTy::mul(j, gsSize);
              auto extractIdx = NodeTy::add(i, jOffset);

              auto stepsNode = NodeTy::constantScalar(steps, DagType::index());
              auto isBound = NodeTy::comparison(extractIdx, stepsNode,
                                                ComparisonPredicate::LT);

              auto buildMultiplyAndAdd = [&]() {
                auto plaintext = extractFunc(babySteppedDag, extractIdx);
                auto innerRotAmount =
                    dagRotationFn(giantStepSize, j, i, period);

                auto rotatedPlaintext =
                    NodeTy::leftRotate(plaintext, innerRotAmount);

                auto babyStepVal = NodeTy::extract(outerIterArgs[1], i);
                auto multiplied = NodeTy::mul(rotatedPlaintext, babyStepVal);
                return NodeTy::yield({NodeTy::add(innerSum, multiplied)});
              };

              auto buildInnerSumYield = [&]() {
                return NodeTy::yield({innerSum});
              };

              NodePtr term;
              if (isNotZeroDiagTensor != nullptr) {
                auto isNotZeroDiag =
                    NodeTy::extract(isNotZeroDiagTensor, extractIdx);
                auto zero = NodeTy::constantScalar(0.0, DagType::floatTy(64));
                auto isNotZero = NodeTy::comparison(isNotZeroDiag, zero,
                                                    ComparisonPredicate::NE);
                term = NodeTy::ifElse(
                    isBound,
                    [&]() {
                      auto innerIf = NodeTy::ifElse(
                          isNotZero, buildMultiplyAndAdd, buildInnerSumYield);
                      return NodeTy::yield({NodeTy::resultAt(innerIf, 0)});
                    },
                    buildInnerSumYield);
              } else {
                term = NodeTy::ifElse(isBound, buildMultiplyAndAdd,
                                      buildInnerSumYield);
              }

              return NodeTy::yield({NodeTy::resultAt(term, 0)});
            });

        auto innerSum = NodeTy::resultAt(innerLoop, 0);

        // Rotate by j * giantStepSize * period
        auto gsSize = NodeTy::constantScalar(giantStepSize, DagType::index());
        auto periodNode = NodeTy::constantScalar(period, DagType::index());
        auto outerRotAmount = NodeTy::mul(j, gsSize);
        outerRotAmount = NodeTy::mul(outerRotAmount, periodNode);

        auto rotatedSum = NodeTy::leftRotate(innerSum, outerRotAmount);

        // Accumulate into outer sum
        auto newOuterSum = NodeTy::add(outerSum, rotatedSum);

        return NodeTy::yield({newOuterSum, outerIterArgs[1]});
      });

  return NodeTy::resultAt(outerLoop, 0);
}

// Returns an arithmetic DAG that implements a tensor_ext.rotate_and_reduce op.
//
// See TensorExtOps.td docs for RotateAndReduceOp for more details.
//
// The `vector` argument is a ciphertext value that will be rotated O(sqrt(n))
// times when the `plaintexts` argument is set (Baby Step Giant Step), or
// O(log(n)) times when the `plaintexts` argument is not set (log-style
// rotate-and-reduce accumulation).
//
// The `plaintexts` argument, when present, represents a vector of pre-packed
// plaintexts that will be rotated and multiplied with the rotated `vector`
// argument in BSGS style.
//
// Note that using this kernel results in places in the pipeline where a
// plaintext type is rotated, but most FHE implementations don't have a
// plaintext rotation operation (it would be wasteful) and instead expect the
// "plaintext rotation" to apply to the cleartext. HEIR has places in the
// pipeline that support this by converting a rotate(encode(cleartext)) to
// encode(rotate(cleartext)).
//
// Note: If this implementation is updated, it's likely that
// rotateAndReduceRotationIndices in RotationUtils.h needs updating, too
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementRotateAndReduce(const T& vector, std::optional<T> plaintexts,
                         int64_t period, int64_t steps, const DagType& dagType,
                         const std::map<int, bool>& zeroDiagonals = {},
                         const std::string& reduceOp = "arith.addi",
                         bool unroll = true) {
  using NodeTy = ArithmeticDagNode<T>;
  auto performReduction = [&](std::shared_ptr<NodeTy> left,
                              std::shared_ptr<NodeTy> right) {
    if (reduceOp == "arith.addi" || reduceOp == "arith.addf") {
      return NodeTy::add(left, right);
    }

    if (reduceOp == "arith.muli" || reduceOp == "arith.mulf") {
      return NodeTy::mul(left, right);
    }

    // Default to add for unknown operations
    return NodeTy::add(left, right);
  };

  if (!plaintexts.has_value()) {
    if (unroll) {
      return implementRotateAndReduceAccumulation<T>(vector, period, steps,
                                                     performReduction);
    }
    return implementRotateAndReduceAccumulationRolled<T>(
        vector, period, steps, performReduction, dagType);
  }

  assert(reduceOp == "arith.addi" ||
         reduceOp == "arith.addf" &&
             "Baby-step-giant-step rotate-and-reduce only supports addition "
             "as the reduction operation");

  if (unroll) {
    // Unrolled version: uses static extraction
    auto extractFunc = [](std::shared_ptr<NodeTy> babySteppedDag,
                          int64_t extractionIndex) {
      return NodeTy::extract(babySteppedDag, extractionIndex);
    };

    return implementBabyStepGiantStep<T>(
        vector, plaintexts.value(), period, steps, dagType, extractFunc,
        zeroDiagonals, defaultDerivedRotationIndexFn);
  }

  // Rolled version: uses dynamic extraction and DAG rotation function
  auto dynamicExtractFunc = [](std::shared_ptr<NodeTy> babySteppedDag,
                               std::shared_ptr<NodeTy> extractionIndex) {
    return NodeTy::extract(babySteppedDag, extractionIndex);
  };

  return implementBabyStepGiantStepRolled<T>(
      vector, plaintexts.value(), period, steps, dagType, dynamicExtractFunc,
      zeroDiagonals, defaultDagDerivedRotationIndexFn<T>);
}

// Returns an arithmetic DAG that implements a dot product kernel.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementDot(const T& lhs, const T& rhs, int64_t steps,
             const DagType& dagType) {
  using NodeTy = ArithmeticDagNode<T>;
  auto mulDag = NodeTy::mul(NodeTy::leaf(lhs), NodeTy::leaf(rhs));
  return implementRotateAndReduceAccumulation<T>(mulDag, /*period=*/1, steps,
                                                 NodeTy::add);
}

// Returns an arithmetic DAG that implements a broadcasted reduce kernel.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementBroadcastedReduce(
    std::shared_ptr<ArithmeticDagNode<T>> vectorDag,
    std::optional<std::shared_ptr<ArithmeticDagNode<T>>> cleanupMaskDag,
    int64_t period, int64_t steps, int64_t numSlots, const DagType& dagType,
    const std::string& reduceOp = "arith.addi", bool unroll = true) {
  using NodeTy = ArithmeticDagNode<T>;
  using NodePtr = std::shared_ptr<NodeTy>;

  DagReducer<T> reduceFunc = [&](NodePtr lhs, NodePtr rhs) {
    if (reduceOp == "arith.addi" || reduceOp == "arith.addf") {
      return NodeTy::add(lhs, rhs);
    }
    if (reduceOp == "arith.muli" || reduceOp == "arith.mulf") {
      return NodeTy::mul(lhs, rhs);
    }
    return NodeTy::add(lhs, rhs);
  };

  NodePtr reduced;
  if (unroll) {
    reduced = implementRotateAndReduceAccumulation<T>(vectorDag, period, steps,
                                                      reduceFunc);
  } else {
    reduced = implementRotateAndReduceAccumulationRolled<T>(
        vectorDag, period, steps, reduceFunc, dagType);
  }

  // Check Natural Replication
  if (steps * period == numSlots) {
    return reduced;
  }

  // Shift to last slots
  int64_t shiftToLast = numSlots - (steps - 1) * period;
  auto shifted = NodeTy::leftRotate(reduced, shiftToLast);

  NodePtr current = shifted;
  // Cleanup Mask
  if (cleanupMaskDag.has_value()) {
    current = NodeTy::mul(current, cleanupMaskDag.value());
  }

  // Replication Tree (Left rotations only)
  for (int64_t rep_shift = 1; rep_shift < steps; rep_shift *= 2) {
    int64_t rotateAmount = rep_shift * period;
    auto rotated = NodeTy::leftRotate(current, rotateAmount);
    current = reduceFunc(current, rotated);
  }

  return current;
}

// Returns an arithmetic DAG that implements a baby-step-giant-step between
// ciphertexts.
//
// This implements equation 21 in 6.2.2 of LKAA25: "Tricycle: Private
// Transformer Inference with Tricyclic Encodings"
// https://eprint.iacr.org/2025/1200
//
// This differs from the above implementRotateAndReduce in that, instead of a
// set of pre-computed plaintexts, both arguments are individual ciphertexts.
// Normally with one ciphertext, the naive approach uses n - 1 rotations that
// BSGS reduces to c sqrt(n) + O(1) rotations, if both inputs are ciphertexts
// then it converts 2n - 2 total rotations to n + c sqrt(n) + O(1) rotations.
// Essentially, the "n to sqrt(n)" redution applies to the `vector` argument
// only, while the `plaintexts` argument still gets n-1 rotations.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementCiphertextCiphertextBabyStepGiantStep(
    const T& giantSteppedOperand, const T& babySteppedOperand, int64_t period,
    int64_t steps, const DagType& dagType,
    DerivedRotationIndexFn derivedRotationIndexFn) {
  using NodeTy = ArithmeticDagNode<T>;

  // Avoid replicating and re-extracting by simulating the extraction step by
  // just returning the single ciphertext.
  auto extractFunc = [](std::shared_ptr<NodeTy> babySteppedDag,
                        int64_t extractionIndex) { return babySteppedDag; };

  return implementBabyStepGiantStep<T>(giantSteppedOperand, babySteppedOperand,
                                       period, steps, dagType, extractFunc, {},
                                       derivedRotationIndexFn);
}

// Returns an arithmetic DAG that implements the Halevi-Shoup matrix
// multiplication algorithm. This implementation uses a rotate-and-reduce
// operation, followed by a summation of partial sums if the matrix is not
// square.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementHaleviShoup(const T& vector, const T& matrix,
                     std::vector<int64_t> originalMatrixShape,
                     const DagType& dagType,
                     std::map<int, bool> zeroDiagonals = {},
                     bool unroll = true) {
  using NodeTy = ArithmeticDagNode<T>;
  using NodePtr = std::shared_ptr<NodeTy>;
  int64_t numRotations = matrix.getShape()[0];

  auto rotateAndReduceResult = implementRotateAndReduce<T>(
      vector, std::optional<T>(matrix), /*period=*/1,
      /*steps=*/numRotations, dagType, zeroDiagonals,
      /*reduceOp=*/"arith.addi",
      /*unroll=*/unroll);

  auto summedShifts = rotateAndReduceResult;

  int64_t matrixNumRows = nextPowerOfTwo(originalMatrixShape[0]);
  int64_t matrixNumCols = nextPowerOfTwo(originalMatrixShape[1]);

  if (matrixNumRows == matrixNumCols) {
    return summedShifts;
  }

  // Post-processing partial-rotate-and-reduce step required for
  // squat-diagonal packing.
  int64_t numShifts = (int64_t)(log2(matrixNumCols) - log2(matrixNumRows));
  if (unroll) {
    int64_t shift = matrixNumCols / 2;
    for (int64_t i = 0; i < numShifts; ++i) {
      auto rotated = NodeTy::leftRotate(summedShifts, shift);
      summedShifts = NodeTy::add(summedShifts, rotated);
      shift /= 2;
    }

    return summedShifts;
  }

  auto shift = NodeTy::constantScalar(matrixNumCols / 2, DagType::index());
  auto loopNode = NodeTy::loop(
      {summedShifts, shift}, {dagType, DagType::index()},
      /*lower=*/0, /*upper=*/numShifts, /*step=*/1,
      [&](NodePtr iv, const std::vector<NodePtr>& iterArgs) {
        auto currentSum = iterArgs[0];
        auto currentShift = iterArgs[1];
        auto rotated = NodeTy::leftRotate(currentSum, currentShift);
        auto newSum = NodeTy::add(currentSum, rotated);
        auto newShift = NodeTy::floorDiv(currentShift, 2);
        return NodeTy::yield({newSum, newShift});
      });
  return NodeTy::resultAt(loopNode, 0);
}

// Returns an arithmetic DAG that implements the bicyclic matrix multiplication
// algorithm.
//
// The input matrices packedA and packedB are assumed to be properly packed to
// meet the conditions for bicyclic multiplication. That is: both matrices are
// zero-padded so that their dimensions are coprime, they are cyclically
// repeated to fill all the slots of the ciphertext, and they are packed
// according to the bicyclic ordering.
//
// This function produces a kernel using roughly n + 2sqrt(n) - 3 rotations
// (for matrix dimensions all order n), by applying the baby-step-giant-step
// method to reduce the number of rotations of packedA.
//
// This implements the BMM-I algorithm from https://eprint.iacr.org/2024/1762
// with modifications from LKAA25 (https://eprint.iacr.org/2025/1200):
//
//  - A simplification of the rotation formula in Sec 5.2.1 (equation 9).
//  - A baby-step-giant-step optimization of the summation below, from Sec
//    6.2.2 (equation 21).
//
// It computes
//
// C = sum_{c=0}^{n-1} rot(A, r1(c)) * rot(B, r2(c))
//
// where
//
//  r1(c) = cm
//  r2(c) = p(cm(p^{-1}) mod n) mod np
//
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementBicyclicMatmul(const T& packedA, const T& packedB, int64_t m,
                        int64_t n, int64_t p, const DagType& dagType) {
  APInt mAPInt = APInt(64, m);
  APInt nAPInt = APInt(64, n);
  APInt pAPInt = APInt(64, p);

  APInt mInvModN = multiplicativeInverse(mAPInt.urem(nAPInt), nAPInt);
  APInt pInvModN = multiplicativeInverse(pAPInt.urem(nAPInt), nAPInt);

  auto derivedRotationIndexFn = [&](int64_t giantStepSize,
                                    int64_t giantStepIndex,
                                    int64_t babyStepIndex, int64_t period) {
    APInt c(64, giantStepIndex * giantStepSize + babyStepIndex);
    APInt mAPInt(64, m);

    // RotY(c) = (p * (c * m * p^{-1} mod n)) mod (n * p)
    APInt rotyInner = (c * mAPInt * pInvModN.getSExtValue()).urem(nAPInt);
    APInt roty = (rotyInner * pAPInt).urem(nAPInt * pAPInt);

    APInt result = roty - APInt(64, period) * APInt(64, giantStepSize) *
                              APInt(64, giantStepIndex);
    // Ensure rotation index is a positive modulo representative [0, np-1].
    result = result.srem(nAPInt * pAPInt);
    if (result.isNegative()) {
      result += nAPInt * pAPInt;
    }
    return result.getSExtValue();
  };

  return implementCiphertextCiphertextBabyStepGiantStep<T>(
      packedA, packedB, /*period=*/m, /*steps=*/n, dagType,
      derivedRotationIndexFn);
}

// Returns an arithmetic DAG that implements the tricyclic batch matrix
// multiplication algorithm (ciphertext-ciphertext). Uses the tricyclic
// rotation formulas from LKAA25 (Tricycle paper) and applies BSGS to
// reduce rotations on the φ(A) side.
//
// The inputs packedA and packedB are expected to be tricyclic encodings
// φ(A) and φ(B) for tensors A ∈ R^{h×m×n} and B ∈ R^{h×n×p}. The function
// applies equation (22) and the ct-ct BSGS decomposition in Section 6.2.2.
//
// Parameters:
//  - packedA: tricyclic-encoded ciphertext for A (φ(A))
//  - packedB: tricyclic-encoded ciphertext for B (φ(B))
//  - h, m, n, p: tricyclic tensor dimensions (h: batch count / heads)
//
// Returns φ(Z) where Z = batch_matmul(A, B) as a DAG node.
template <typename T>
std::enable_if_t<std::is_base_of<AbstractValue, T>::value,
                 std::shared_ptr<ArithmeticDagNode<T>>>
implementTricyclicBatchMatmul(const T& packedA, const T& packedB, int64_t h,
                              int64_t m, int64_t n, int64_t p,
                              const DagType& dagType) {
  APInt hAPInt = APInt(64, h);
  APInt mAPInt = APInt(64, m);
  APInt nAPInt = APInt(64, n);
  APInt pAPInt = APInt(64, p);

  APInt pInvModN = multiplicativeInverse(pAPInt.urem(nAPInt), nAPInt);

  APInt modulus = (hAPInt * nAPInt * pAPInt);
  // This follows Eq. (22) and the ct-ct BSGS decomposition.
  // RotY(c) = (h * p * ( (c * m * p^{-1}) mod n )) mod (h * n * p)
  auto derivedRotationIndexFn = [&](int64_t giantStepSize,
                                    int64_t giantStepIndex,
                                    int64_t babyStepIndex, int64_t period) {
    APInt c(64, giantStepIndex * giantStepSize + babyStepIndex);

    // rotyInner = (c * m * p^{-1}) mod n
    APInt rotyInner = (c * mAPInt * pInvModN.getSExtValue()).urem(nAPInt);

    // rotY calculation from LKAA25 Eq. (22):
    // RotY(c) = (h * p * (c * m * p^{-1} mod n)) mod (h * n * p)
    APInt roty = (rotyInner * hAPInt * pAPInt).urem(modulus);

    APInt result = roty - APInt(64, period) * APInt(64, giantStepSize) *
                              APInt(64, giantStepIndex);
    result = result.srem(modulus);
    if (result.isNegative()) {
      result += modulus;
    }
    return result.getSExtValue();
  };

  int64_t period = h * m;
  return implementCiphertextCiphertextBabyStepGiantStep<T>(
      packedA, packedB, /*period=*/period, /*steps=*/n, dagType,
      derivedRotationIndexFn);
}

}  // namespace kernel
}  // namespace heir
}  // namespace mlir

#endif  // LIB_KERNEL_KERNELIMPLEMENTATION_H_
