#include <cassert>
#include <cstdlib>
#include <vector>

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

using namespace onnx_mlir::spatial;

int main() {
  TransferCost transfer {.fixed = 50, .networkFlits = 4};

  SchedulingTarget fast;
  fast.processorCount = 2;
  fast.interProcessorLatencyNs = {0, 3, 3, 0};
  fast.averageInterProcessorLatencyNs = 3;

  SchedulingTarget slow = fast;
  slow.interProcessorLatencyNs = {0, 10, 10, 0};
  slow.averageInterProcessorLatencyNs = 10;

  assert(getPeftTransferTime(transfer, 0, 0, fast) == 0);
  assert(getPeftTransferTime(transfer, 0, 1, fast) == 62);
  assert(getPeftTransferTime(transfer, 0, 1, slow) == 90);
  assert(fast.getInterProcessorLatencyNs(0, 1) == 3);
  assert(slow.getInterProcessorLatencyNs(0, 1) == 10);

  SchedulingTarget line;
  line.processorCount = 3;
  line.interProcessorLatencyNs = {
    0,
    1,
    10,
    1,
    0,
    1,
    10,
    1,
    0,
  };
  std::vector<Cost> logicalTrafficFlits(9, 0);
  logicalTrafficFlits[2] = 100;
  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, line) == std::vector<size_t>({1, 0, 2}));

  SchedulingTarget alreadyPlaced = line;
  alreadyPlaced.interProcessorLatencyNs = {
    0,
    10,
    1,
    10,
    0,
    1,
    1,
    1,
    0,
  };
  assert(mapLogicalProcessorsToPhysicalCores(logicalTrafficFlits, alreadyPlaced) == std::vector<size_t>({0, 1, 2}));
  return EXIT_SUCCESS;
}