#ifndef LIB_UTILS_LAYOUT_EVALUATE_H_
#define LIB_UTILS_LAYOUT_EVALUATE_H_

#include <cassert>
#include <cstdint>
#include <functional>
#include <optional>
#include <vector>

#include "lib/Utils/Layout/IslConversion.h"
#include "lib/Utils/Layout/Utils.h"
#include "mlir/include/mlir/Analysis/Presburger/IntegerRelation.h"  // from @llvm-project
#include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h"  // from @llvm-project

// ISL
#include "include/isl/ctx.h"                 // from @isl
#include "include/isl/map.h"                 // from @isl
#include "include/isl/map_type.h"            // from @isl
#include "include/isl/point.h"               // from @isl
#include "include/isl/set.h"                 // from @isl
#include "include/isl/space_type.h"          // from @isl
#include "include/isl/val.h"                 // from @isl
#include "include/isl/val_type.h"            // from @isl
#include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace heir {

using presburger::BoundType;
using presburger::VarKind;

template <typename T>
std::vector<std::vector<T>> evaluateLayout(
    const presburger::IntegerRelation& relation,
    std::function<T(const std::vector<int64_t>&)> getValueFn,
    std::optional<SmallVector<int64_t>> outputShape = std::nullopt) {
  int64_t numCt, numSlot;
  if (outputShape.has_value()) {
    numCt = outputShape.value()[0];
    numSlot = outputShape.value()[1];
  } else {
    auto ctBound = relation.getConstantBound64(
        BoundType::UB, relation.getVarKindOffset(VarKind::Range));
    assert(ctBound.has_value() &&
           "expected a constant bound for the number of ciphertexts");
    auto slotBond = relation.getConstantBound64(
        BoundType::UB, relation.getVarKindOffset(VarKind::Range) + 1);
    assert(slotBond.has_value() &&
           "expected a constant bound for the number of slots");
    numCt = ctBound.value() + 1;
    numSlot = slotBond.value() + 1;
  }
  std::vector<std::vector<T>> result(numCt, std::vector<T>(numSlot, 0));

  // Get all points in the relation.
  PointPairCollector collector(relation.getNumDomainVars(), /*rangeDims=*/2);
  enumeratePoints(relation, collector);

  for (const auto& pointPair : collector.points) {
    std::vector<int64_t> domainPoint = pointPair.first;
    std::vector<int64_t> rangePoint = pointPair.second;
    result[rangePoint[0]][rangePoint[1]] = getValueFn(domainPoint);
  }
  return result;
}

// This applies the layout relation on a given input vector.
//
// The layout relation is a presburger relation with domain size equal to the
// rank of the input vector and range size equal to two (ct, slot). The return
// value is a 2-D vector whose ct, slot entries consist of the data that is
// mapped from the layout relation.

// This is intended to be used for test vectors, since the performance is not
// optimized.
template <typename T>
std::vector<std::vector<T>> evaluateLayoutOnVector(
    const presburger::IntegerRelation& relation, const std::vector<T>& input,
    std::optional<SmallVector<int64_t>> outputShape = std::nullopt) {
  std::function<T(const std::vector<int64_t>&)> getValueFn =
      [&](const std::vector<int64_t>& domainPoint) {
        return input[domainPoint[0]];
      };
  return evaluateLayout(relation, getValueFn, outputShape);
}

// This applies the layout relation on a given input 2-D vector. See above.

template <typename T>
std::vector<std::vector<T>> evaluateLayoutOnMatrix(
    const presburger::IntegerRelation& relation,
    const std::vector<std::vector<T>>& input) {
  std::function<T(const std::vector<int64_t>&)> getValueFn =
      [&](const std::vector<int64_t>& domainPoint) {
        return input[domainPoint[0]][domainPoint[1]];
      };
  return evaluateLayout(relation, getValueFn);
}

// Evaluate a layout relation on a 3-D tensor (domain ordering: h, m, n).
// Returns a vector of ciphertexts, each ciphertext is a vector of slots.
template <typename T>
std::vector<std::vector<T>> evaluateLayoutOnTensor(
    const presburger::IntegerRelation& relation,
    const std::vector<std::vector<std::vector<T>>>& tensor) {
  std::function<T(const std::vector<int64_t>&)> getValueFn =
      [&](const std::vector<int64_t>& domainPoint) -> T {
    return tensor[domainPoint[0]][domainPoint[1]][domainPoint[2]];
  };
  return evaluateLayout(relation, getValueFn);
}

// This unpacks a layout relation (ct, slot) on a given n-D vector.
//
// This is intended to be used for test vectors, since the performance is not
// optimized.
template <typename T>
void unpackLayout(
    const presburger::IntegerRelation& relation,
    const std::vector<std::vector<T>>& packed,
    std::function<void(const std::vector<int64_t>&, T)> setValueFn) {
  // Get all points in the relation.
  PointPairCollector collector(relation.getNumDomainVars(), /*rangeDims=*/2);
  enumeratePoints(relation, collector);

  for (const auto& pointPair : collector.points) {
    const std::vector<int64_t>& domainPoint = pointPair.first;
    const std::vector<int64_t>& rangePoint = pointPair.second;
    setValueFn(domainPoint, packed[rangePoint[0]][rangePoint[1]]);
  }
}

// Unpack a packed (ct x slots) representation back to a 1-D vector with the
// given originalShape {n}.
template <typename T>
std::vector<T> unpackLayoutToVector(const presburger::IntegerRelation& relation,
                                    const std::vector<std::vector<T>>& packed,
                                    ArrayRef<int64_t> originalShape) {
  assert(originalShape.size() == 1 && "originalShape must be 1-D");
  std::vector<T> result(originalShape[0], 0);

  unpackLayout<T>(relation, packed,
                  [&](const std::vector<int64_t>& domainPoint, T value) {
                    result[domainPoint[0]] = value;
                  });
  return result;
}

// Unpack a packed (ct x slots) representation back to a 2-D matrix with the
// given originalShape {m, n}.
template <typename T>
std::vector<std::vector<T>> unpackLayoutToMatrix(
    const presburger::IntegerRelation& relation,
    const std::vector<std::vector<T>>& packed,
    ArrayRef<int64_t> originalShape) {
  assert(originalShape.size() == 2 && "originalShape must be 2-D");
  std::vector<std::vector<T>> result(originalShape[0],
                                     std::vector<T>(originalShape[1], 0));

  unpackLayout<T>(relation, packed,
                  [&](const std::vector<int64_t>& domainPoint, T value) {
                    result[domainPoint[0]][domainPoint[1]] = value;
                  });
  return result;
}

// Unpack a packed (ct x slots) representation back to a 3-D tensor with the
// given originalShape {c, m, n}.
template <typename T>
std::vector<std::vector<std::vector<T>>> unpackLayoutTo3DTensor(
    const presburger::IntegerRelation& relation,
    const std::vector<std::vector<T>>& packed,
    ArrayRef<int64_t> originalShape) {
  assert(originalShape.size() == 3 && "originalShape must be 3-D");
  std::vector<std::vector<std::vector<T>>> result(
      originalShape[0],
      std::vector<std::vector<T>>(originalShape[1],
                                  std::vector<T>(originalShape[2], 0)));

  unpackLayout<T>(
      relation, packed, [&](const std::vector<int64_t>& domainPoint, T value) {
        result[domainPoint[0]][domainPoint[1]][domainPoint[2]] = value;
      });
  return result;
}

// Unpack a packed (ct x slots) representation back to a 4-D tensor with the
// given originalShape {c, h, m, n}.
template <typename T>
std::vector<std::vector<std::vector<std::vector<T>>>> unpackLayoutTo4DTensor(
    const presburger::IntegerRelation& relation,
    const std::vector<std::vector<T>>& packed,
    ArrayRef<int64_t> originalShape) {
  assert(originalShape.size() == 4 && "originalShape must be 4-D");
  std::vector<std::vector<std::vector<std::vector<T>>>> result(
      originalShape[0],
      std::vector<std::vector<std::vector<T>>>(
          originalShape[1],
          std::vector<std::vector<T>>(originalShape[2],
                                      std::vector<T>(originalShape[3], 0))));

  unpackLayout<T>(
      relation, packed, [&](const std::vector<int64_t>& domainPoint, T value) {
        result[domainPoint[0]][domainPoint[1]][domainPoint[2]][domainPoint[3]] =
            value;
      });
  return result;
}

}  // namespace heir
}  // namespace mlir

#endif  // LIB_UTILS_LAYOUT_EVALUATE_H_
