#include <cstdint>
#include <memory>
#include <optional>
#include <vector>

#include "gtest/gtest.h"  // from @googletest
#include "lib/Kernel/AbstractValue.h"
#include "lib/Kernel/ArithmeticDag.h"
#include "lib/Kernel/EvalVisitor.h"
#include "lib/Kernel/KernelImplementation.h"

namespace mlir {
namespace heir {
namespace kernel {
namespace {

// A `tensor_ext.rotate_and_reduce` op computes the reduction of rotated tensors
// in the form `\sum_{i = 0}^{n} P(i) \cdot rotate(v, T*i)` where `T` is some
// period of rotation. The naive approach would compute `n` rotations of the
// ciphertext `v`.

std::vector<int> rotate(const std::vector<int>& vec, int64_t amount) {
  int64_t n = vec.size();
  std::vector<int> result(n);
  for (int64_t i = 0; i < n; ++i) {
    // rotate left cyclically
    result[i] = vec[(i + amount) % n];
  }
  return result;
}

TEST(RotateAndReduceImplTest, TestRotateHelper) {
  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<int> expected = {2, 3, 4, 5, 6, 7, 0, 1};
  std::vector<int> actual = rotate(vector, 2);
  EXPECT_EQ(expected, actual);
}

std::vector<int> runNaive(const std::vector<int>& vec,
                          const std::vector<std::vector<int>>& plaintexts,
                          int64_t period, int64_t n) {
  std::vector<int> result(vec.size(), 0);
  for (int64_t i = 0; i < n; ++i) {
    std::vector<int> rotated = rotate(vec, period * i);
    for (int64_t j = 0; j < vec.size(); ++j) {
      result[j] += plaintexts[i][j] * rotated[j];
    }
  }
  return result;
}

std::vector<int> runImpl(const std::vector<int>& vec,
                         const std::vector<std::vector<int>>& plaintexts,
                         int64_t period, int64_t n, bool unroll = true) {
  LiteralValue vectorInput(vec);

  std::shared_ptr<ArithmeticDagNode<LiteralValue>> result;
  std::optional<LiteralValue> plaintextsInput = std::nullopt;
  if (!plaintexts.empty()) {
    plaintextsInput = std::optional<LiteralValue>(LiteralValue(plaintexts));
  }
  result = implementRotateAndReduce(
      vectorInput, plaintextsInput, period, n,
      DagType::intTensor(32, {static_cast<int64_t>(vec.size())}), {},
      "arith.addi", unroll);
  return std::get<std::vector<int>>(evalKernel(result)[0].get());
}

TEST(RotateAndReduceImplTest, TestUnitPeriodWithPlaintextsSimpleValues) {
  int64_t n = 8;
  int64_t period = 1;

  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::vector<int>> plaintexts = {
      {1, 0, 0, 0, 0, 0, 0, 0}, {0, 0, 0, 0, 0, 0, 0, 0},
      {0, 0, 0, 0, 0, 0, 0, 0}, {0, 0, 0, 0, 0, 0, 0, 0},
      {0, 0, 0, 0, 0, 0, 0, 0}, {0, 0, 0, 0, 0, 0, 0, 0},
      {0, 0, 0, 0, 0, 0, 0, 0}, {0, 0, 0, 0, 0, 0, 0, 0},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestUnitPeriodWithPlaintextsSimpleValuesRolled) {
  int64_t n = 4;
  int64_t period = 1;

  std::vector<int> vector = {0, 1, 2, 3};
  std::vector<std::vector<int>> plaintexts = {
      {1, 0, 0, 0},
      {0, 0, 0, 0},
      {0, 0, 0, 0},
      {0, 0, 0, 0},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n, false);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestUnitPeriodWithPlaintexts) {
  int64_t n = 8;
  int64_t period = 1;

  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::vector<int>> plaintexts = {
      {1, 2, 3, 4, 5, 6, 7, 8},      {2, 3, 4, 5, 6, 7, 8, 9},
      {3, 4, 5, 6, 7, 8, 9, 10},     {4, 5, 6, 7, 8, 9, 10, 11},
      {5, 6, 7, 8, 9, 10, 11, 12},   {6, 7, 8, 9, 10, 11, 12, 13},
      {7, 8, 9, 10, 11, 12, 13, 14}, {8, 9, 10, 11, 12, 13, 14, 15},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestPeriod2WithPlaintexts) {
  int64_t n = 8;
  int64_t period = 2;

  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::vector<int>> plaintexts = {
      {1, 2, 3, 4, 5, 6, 7, 8},      {2, 3, 4, 5, 6, 7, 8, 9},
      {3, 4, 5, 6, 7, 8, 9, 10},     {4, 5, 6, 7, 8, 9, 10, 11},
      {5, 6, 7, 8, 9, 10, 11, 12},   {6, 7, 8, 9, 10, 11, 12, 13},
      {7, 8, 9, 10, 11, 12, 13, 14}, {8, 9, 10, 11, 12, 13, 14, 15},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestPeriod2WithPlaintextsSmallerN) {
  int64_t n = 4;
  int64_t period = 2;

  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<std::vector<int>> plaintexts = {
      {1, 2, 3, 4, 5, 6, 7, 8},
      {2, 3, 4, 5, 6, 7, 8, 9},
      {3, 4, 5, 6, 7, 8, 9, 10},
      {4, 5, 6, 7, 8, 9, 10, 11},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestPeriod1WithNoPlaintext) {
  int64_t n = 8;
  int64_t period = 1;
  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<int> expected(8, 28);
  std::vector<int> actual = runImpl(vector, {}, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestPeriod1WithNoPlaintextRolled) {
  int64_t n = 8;
  int64_t period = 1;
  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<int> expected(8, 28);
  std::vector<int> actual = runImpl(vector, {}, period, n, false);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestPeriod2WithNoPlaintext) {
  int64_t n = 4;
  int64_t period = 2;
  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  // This test computes the partial sum of entries with a stride of 2, i.e.,
  // every even index contains the sum of all even indices and every odd index
  // contains the sum of all odd indices.
  std::vector<int> expected = {12, 16, 12, 16, 12, 16, 12, 16};
  std::vector<int> actual = runImpl(vector, {}, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestUnitPeriodWithIndivisibleN10) {
  int64_t n = 10;
  int64_t period = 1;

  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  std::vector<std::vector<int>> plaintexts = {
      {1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
      {2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
      {3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
      {4, 5, 6, 7, 8, 9, 10, 11, 12, 13},
      {5, 6, 7, 8, 9, 10, 11, 12, 13, 14},
      {6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
      {7, 8, 9, 10, 11, 12, 13, 14, 15, 16},
      {8, 9, 10, 11, 12, 13, 14, 15, 16, 17},
      {9, 10, 11, 12, 13, 14, 15, 16, 17, 18},
      {10, 11, 12, 13, 14, 15, 16, 17, 18, 19},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, TestUnitPeriodWithIndivisibleN6) {
  int64_t n = 6;
  int64_t period = 1;

  std::vector<int> vector = {0, 1, 2, 3, 4, 5};
  std::vector<std::vector<int>> plaintexts = {
      {1, 2, 3, 4, 5, 6}, {2, 3, 4, 5, 6, 7},  {3, 4, 5, 6, 7, 8},
      {4, 5, 6, 7, 8, 9}, {5, 6, 7, 8, 9, 10}, {6, 7, 8, 9, 10, 11},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
  EXPECT_EQ(expected, actual);
}

TEST(RotateAndReduceImplTest, RegressionTest) {
  int64_t n = 4;
  int64_t period = 1;

  std::vector<int> vector = {0, 1, 2, 3, 4, 0, 0, 0, 0, 1, 2, 3, 4, 0, 0, 0,
                             0, 1, 2, 3, 4, 0, 0, 0, 0, 1, 2, 3, 4, 0, 0, 0};
  std::vector<std::vector<int>> plaintexts = {
      {1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0,
       1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0},
      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
  };

  std::vector<int> expected = runNaive(vector, plaintexts, period, n);
  std::vector<int> actual = runImpl(vector, plaintexts, period, n);
  EXPECT_EQ(expected, actual);
}

std::vector<int> runBroadcastedReduceImpl(
    const std::vector<int>& vec, std::optional<std::vector<int>> cleanupMask,
    int64_t period, int64_t steps, bool unroll = true) {
  using NodeTy = ArithmeticDagNode<LiteralValue>;
  using NodePtr = std::shared_ptr<NodeTy>;

  LiteralValue vectorInput(vec);
  auto vectorDag = NodeTy::leaf(vectorInput);

  std::optional<NodePtr> cleanupMaskDag = std::nullopt;
  if (cleanupMask.has_value()) {
    cleanupMaskDag = NodeTy::leaf(LiteralValue(cleanupMask.value()));
  }

  auto result = implementBroadcastedReduce<LiteralValue>(
      vectorDag, cleanupMaskDag, period, steps, vec.size(),
      DagType::intTensor(32, {static_cast<int64_t>(vec.size())}), "arith.addi",
      unroll);

  return std::get<std::vector<int>>(evalKernel(result)[0].get());
}

TEST(RotateAndReduceImplTest, BroadcastedReduce_Natural_PowerOfTwo) {
  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<int> expected(8, 28);

  for (bool unroll : {true, false}) {
    std::vector<int> actual =
        runBroadcastedReduceImpl(vector, std::nullopt, 1, 8, unroll);
    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
  }
}

TEST(RotateAndReduceImplTest, BroadcastedReduce_Natural_Stride) {
  std::vector<int> vector = {0, 1, 2, 3, 4, 5, 6, 7};
  std::vector<int> expected = {12, 16, 12, 16, 12, 16, 12, 16};

  for (bool unroll : {true, false}) {
    std::vector<int> actual =
        runBroadcastedReduceImpl(vector, std::nullopt, 2, 4, unroll);
    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
  }
}

TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Contiguous) {
  std::vector<int> vector = {0,  1,  2,  3,  4,  5,  6,  7,
                             10, 11, 12, 13, 14, 15, 16, 17};
  std::vector<int> mask = {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1};
  std::vector<int> expected = {28,  28,  28,  28,  28,  28,  28,  28,
                               108, 108, 108, 108, 108, 108, 108, 108};

  for (bool unroll : {true, false}) {
    std::vector<int> actual =
        runBroadcastedReduceImpl(vector, mask, 1, 8, unroll);
    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
  }
}

TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Stride) {
  std::vector<int> vector = {0,  1,  2,  3,  4,  5,  6,  7,
                             10, 11, 12, 13, 14, 15, 16, 17};
  std::vector<int> mask = {0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1};
  std::vector<int> expected = {12, 16, 12, 16, 12, 16, 12, 16,
                               52, 56, 52, 56, 52, 56, 52, 56};

  for (bool unroll : {true, false}) {
    std::vector<int> actual =
        runBroadcastedReduceImpl(vector, mask, 2, 4, unroll);
    EXPECT_EQ(expected, actual) << "Failed for unroll=" << unroll;
  }
}

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