#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp"

#include <cassert>
#ifdef NDEBUG
#error "Periodic dependency tests require assertions"
#endif
#include <cmath>
#include <functional>
#include <random>

using namespace onnx_mlir::spatial;

static double enumerateCycles(size_t count,
                             const std::vector<PeriodicDependency> &edges) {
  double maximum = 0;
  std::vector<bool> visited(count);
  for (size_t start = 0; start < count; ++start) {
    std::function<void(size_t, Cost, Cost)> visit =
      [&](size_t source, Cost delay, Cost distance) {
        visited[source] = true;
        for (const auto &edge : edges) {
          if (edge.source != source)
            continue;
          Cost nextDelay = delay + edge.delayNs;
          Cost nextDistance = distance + edge.distance;
          if (edge.target == start && nextDistance)
            maximum = std::max(maximum,
              static_cast<double>(nextDelay) / nextDistance);
          else if (!visited[edge.target])
            visit(edge.target, nextDelay, nextDistance);
        }
        visited[source] = false;
      };
    visit(start, 0, 0);
  }
  return maximum;
}

static void check(size_t count, const std::vector<PeriodicDependency> &edges,
                  double expected) {
  std::string error;
  auto result = analyzePeriodicDependencies(count, edges, error);
  assert(mlir::succeeded(result) && error.empty());
  assert(result->recurrenceLowerNs <= expected + 1e-6);
  assert(result->recurrenceUpperNs >= expected - 1e-6);
  assert(result->recurrenceUpperNs - result->recurrenceLowerNs <= 1.000001);
  Cost delay = 0, distance = 0;
  for (size_t index : result->criticalCycle) {
    delay += edges[index].delayNs;
    distance += edges[index].distance;
  }
  if (expected > 1) {
    assert(distance > 0);
    assert(static_cast<double>(delay) / distance >= result->recurrenceLowerNs - 1e-6);
  }
  auto repeated = analyzePeriodicDependencies(count, edges, error);
  assert(repeated->criticalCycle == result->criticalCycle);
  assert(repeated->recurrenceUpperNs == result->recurrenceUpperNs);
}

int main() {
  check(3, {{0, 1, 12, 0}, {1, 2, 7, 0}}, 0);
  check(3, {{0, 1, 12, 1}, {1, 2, 7, 1}}, 0);
  check(1, {{0, 0, 30, 2}}, 15);
  // Independent core periods remain parallel.
  check(2, {{0, 0, 10, 1}, {1, 1, 20, 1}}, 20);
  // Equal serial demand, different cross-core blocking recurrences.
  check(4, {{0, 1, 10, 0}, {2, 3, 10, 0},
            {1, 0, 0, 1}, {3, 2, 0, 1}}, 10);
  check(4, {{0, 1, 10, 0}, {2, 3, 10, 0},
            {1, 2, 0, 0}, {3, 0, 0, 1}}, 20);
  std::string error;
  assert(mlir::failed(analyzePeriodicDependencies(
    2, {{0, 1, 1, 0}, {1, 0, 1, 0}}, error)));
  assert(mlir::failed(analyzePeriodicDependencies(1, {{0, 1, 1, 1}}, error)));
  assert(mlir::failed(analyzePeriodicDependencies(
    1, {{0, 0, Cost{1} << 60, 1}}, error)));
  std::mt19937 random(812);
  for (unsigned sample = 0; sample < 300; ++sample) {
    std::vector<PeriodicDependency> edges;
    for (size_t source = 0; source < 6; ++source)
      for (size_t target = 0; target < 6; ++target)
        if (random() % 4 == 0)
          edges.push_back({source, target, random() % 101,
            source < target ? 0u : static_cast<unsigned>(1u + random() % 3)});
    check(6, edges, enumerateCycles(6, edges));
  }
}