#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 <chrono>
#include <iostream>
#include <functional>
#include <random>

using namespace onnx_mlir::spatial;

namespace onnx_mlir::spatial {
mlir::FailureOr<PeriodicDependencyEstimate> analyzeFrozenPeriodicDependencies(
    size_t, llvm::ArrayRef<PeriodicDependency>, std::string &);
}

static size_t differentialCases = 0, acceleratedRounds = 0, frozenRounds = 0;
static double acceleratedSeconds = 0, frozenSeconds = 0;

static void differential(size_t count, const std::vector<PeriodicDependency> &edges) {
  std::string actualError, referenceError;
  auto start = std::chrono::steady_clock::now();
  auto actual = analyzePeriodicDependencies(count, edges, actualError);
  acceleratedSeconds += std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
  start = std::chrono::steady_clock::now();
  auto reference = analyzeFrozenPeriodicDependencies(count, edges, referenceError);
  frozenSeconds += std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
  ++differentialCases;
  assert(mlir::succeeded(actual) == mlir::succeeded(reference));
  assert(actualError == referenceError);
  if (mlir::failed(actual)) return;
  acceleratedRounds += actual->relaxationRounds;
  frozenRounds += reference->relaxationRounds;
  assert(actual->recurrenceLowerNs == reference->recurrenceLowerNs);
  assert(actual->recurrenceUpperNs == reference->recurrenceUpperNs);
  assert(actual->withinGenerationSpanNs == reference->withinGenerationSpanNs);
  assert(actual->criticalCycle == reference->criticalCycle);
}

static void testFrozenDifferential() {
  std::mt19937 random(171221);
  for (size_t sample = 0; sample < 800; ++sample) {
    const size_t count = sample < 600 ? 2 + random() % 31 : 64 + random() % 193;
    std::vector<PeriodicDependency> edges;
    for (size_t source = 0; source < count; ++source) {
      const size_t component = source / 32;
      for (size_t link = 0; link < 4; ++link) {
        size_t target = random() % count;
        if (sample % 3 == 0 && target / 32 != component) continue;
        unsigned distance = target > source && link % 2 == 0 ? 0 : 1 + random() % 7;
        edges.push_back({source, target, random() % 100001, distance});
      }
    }
    differential(count, edges);
  }
  for (size_t sample = 0; sample < 100; ++sample) {
    std::vector<PeriodicDependency> edges;
    for (size_t node = 0; node < 32; ++node) {
      edges.push_back({node, node, (Cost{1} << 40) + random() % 10001,
        static_cast<unsigned>(1024 + random() % 7)});
      edges.push_back({node, (node + 1) % 32, random() % 101, 1});
    }
    differential(32, edges);
  }
  differential(2, {{0,1,1,0}, {1,0,1000000000000ULL, std::numeric_limits<unsigned>::max()}});
  for (size_t count : {2u, 64u, 1024u}) {
    std::vector<PeriodicDependency> edges;
    for (size_t node = 0; node < count; ++node) {
      edges.push_back({node, node, 1000000 + node % 2, 3});
      if (node + 1 < count) edges.push_back({node, node + 1, 1, 0});
    }
    differential(count, edges);
  }
  std::vector<PeriodicDependency> large;
  constexpr size_t componentSize = 2048, components = 32;
  for (size_t component = 0; component < components; ++component) {
    size_t first = component * componentSize;
    for (size_t node = first; node + 1 < first + componentSize; ++node)
      large.push_back({node, node + 1, 1, 0});
    large.push_back({first + componentSize - 1, first, 17 + component,
      static_cast<unsigned>(1 + component % 7)});
  }
  differential(componentSize * components, large);
  differential(0, {});
  differential(2, {{0, 1, 0, 0}, {1, 0, 0, 0}});
  differential(2, {{0, 1, 10, 0}, {1, 0, 3, 0}});
  differential(1, {{0, 1, 1, 1}});
  differential(1, {{0, 0, Cost{1} << 60, 1}});
  differential(1, {{0, 0, std::numeric_limits<Cost>::max(), 1}, {0, 0, 1, 1}});
  std::cout << "frozen_differential_cases=" << differentialCases
    << " accelerated_rounds=" << acceleratedRounds << " frozen_rounds=" << frozenRounds
    << " accelerated_seconds=" << acceleratedSeconds << " frozen_seconds=" << frozenSeconds << '\n';
}

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() {
  testFrozenDifferential();
  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);
  // Positive-cycle checks must not change the optimum or the final witness.
  std::vector<PeriodicDependency> independent;
  for (size_t core = 0; core < 128; ++core)
    independent.push_back({core, core, 1000 + core, 1});
  check(128, independent, 1127);
  std::string performanceError;
  auto accelerated = analyzePeriodicDependencies(128, independent, performanceError);
  assert(mlir::succeeded(accelerated));
  assert(accelerated->criticalCycle == std::vector<size_t>{127});
  assert(accelerated->relaxationRounds < 512);
  // 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));
  }
}