// Test-only reference from commit 3c7679594beb29ba9497aaf3e341cee74fa0be50.
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeriodicDependencies.hpp"

#include <limits>

namespace onnx_mlir::spatial {

mlir::FailureOr<PeriodicDependencyEstimate> analyzeFrozenPeriodicDependencies(
    size_t nodeCount, llvm::ArrayRef<PeriodicDependency> edges,
    std::string &error) {
  std::vector<std::vector<size_t>> outgoing(nodeCount);
  std::vector<size_t> indegree(nodeCount), recurrent;
  Cost upper = 0;
  for (auto [index, edge] : llvm::enumerate(edges)) {
    if (edge.source >= nodeCount || edge.target >= nodeCount) {
      error = "periodic dependency references an invalid event";
      return mlir::failure();
    }
    if (edge.delayNs > std::numeric_limits<Cost>::max() - upper) {
      error = "periodic delay sum exceeds the timing range";
      return mlir::failure();
    }
    upper += edge.delayNs;
    if (edge.distance) {
      recurrent.push_back(index);
    } else {
      outgoing[edge.source].push_back(index);
      ++indegree[edge.target];
    }
  }
  std::vector<size_t> order;
  for (size_t node = 0; node < nodeCount; ++node)
    if (!indegree[node])
      order.push_back(node);
  for (size_t cursor = 0; cursor < order.size(); ++cursor)
    for (size_t index : outgoing[order[cursor]])
      if (!--indegree[edges[index].target])
        order.push_back(edges[index].target);
  if (order.size() != nodeCount) {
    error = "periodic distance-zero dependencies contain a cycle";
    return mlir::failure();
  }

  PeriodicDependencyEstimate result;
  std::vector<Cost> span(nodeCount);
  for (size_t node : order) {
    result.withinGenerationSpanNs =
      std::max(result.withinGenerationSpanNs, span[node]);
    for (size_t index : outgoing[node]) {
      const auto &edge = edges[index];
      span[edge.target] = std::max(span[edge.target],
        checkedAdd(span[node], edge.delayNs));
    }
  }
  if (recurrent.empty() || !upper)
    return result;

  const size_t absent = std::numeric_limits<size_t>::max();
  std::vector<double> potential(nodeCount);
  std::vector<size_t> predecessor(nodeCount, absent);
  auto infeasible = [&](double interval) {
    std::fill(potential.begin(), potential.end(), 0);
    std::fill(predecessor.begin(), predecessor.end(), absent);
    auto relax = [&](size_t index) {
      const auto &edge = edges[index];
      double value = potential[edge.source] + edge.delayNs
                   - interval * edge.distance;
      if (value <= potential[edge.target] + 1e-7)
        return false;
      potential[edge.target] = value;
      predecessor[edge.target] = index;
      return true;
    };
    for (size_t node : order)
      for (size_t index : outgoing[node])
        relax(index);
    // Eliminate distance-zero paths by a DAG sweep. A simple residual path
    // visits at most R recurrent edges, so R+1 rounds detect a positive cycle.
    for (size_t round = 0; round <= recurrent.size(); ++round) {
      ++result.relaxationRounds;
      bool changed = false;
      for (size_t index : recurrent)
        changed |= relax(index);
      for (size_t node : order)
        for (size_t index : outgoing[node])
          changed |= relax(index);
      if (!changed)
        return false;
    }
    return true;
  };
  result.recurrenceUpperNs = upper;
  for (unsigned iteration = 0;
       iteration < 64 && result.recurrenceUpperNs - result.recurrenceLowerNs > 1;
       ++iteration) {
    double middle = (result.recurrenceLowerNs + result.recurrenceUpperNs) / 2;
    if (middle == result.recurrenceLowerNs || middle == result.recurrenceUpperNs)
      break;
    if (infeasible(middle))
      result.recurrenceLowerNs = middle;
    else
      result.recurrenceUpperNs = middle;
  }
  if (result.recurrenceUpperNs - result.recurrenceLowerNs > 1) {
    error = "periodic timing exceeds one-nanosecond floating-point precision";
    return mlir::failure();
  }
  if (!infeasible(result.recurrenceLowerNs))
    return result;

  // Recover witness cycles from the last longest-path predecessor forest.
  std::vector<size_t> seen(nodeCount, absent);
  double best = 0;
  for (size_t start = 0; start < nodeCount; ++start) {
    size_t node = start;
    while (seen[node] == absent && predecessor[node] != absent) {
      seen[node] = start;
      node = edges[predecessor[node]].source;
    }
    if (seen[node] != start || predecessor[node] == absent)
      continue;
    size_t first = node;
    Cost delay = 0, distance = 0;
    std::vector<size_t> cycle;
    do {
      size_t index = predecessor[node];
      cycle.push_back(index);
      delay = checkedAdd(delay, edges[index].delayNs);
      distance = checkedAdd(distance, static_cast<Cost>(edges[index].distance));
      node = edges[index].source;
    } while (node != first);
    double ratio = distance ? static_cast<double>(delay) / distance : 0;
    if (ratio > best) {
      best = ratio;
      std::reverse(cycle.begin(), cycle.end());
      result.criticalCycle = std::move(cycle);
    }
  }
  return result;
}

} // namespace onnx_mlir::spatial