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

#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.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}));

  ComputeGraph graph;
  graph.successors.resize(6);
  graph.predecessors.resize(6);
  graph.successors[1].push_back({2, TransferCost {.fixed = 1, .networkFlits = 1}});
  graph.predecessors[2].push_back({1, TransferCost {.fixed = 1, .networkFlits = 1}});
  const Cost costs[] = {6, 4, 6, 4, 8, 8};
  for (uint32_t task = 0; task < 6; ++task) {
    ComputeInstance instance {nullptr, task, 1};
    ResidentWeight weight;
    weight.opaqueLane = task;
    graph.nodes.push_back({instance, costs[task], {weight}, task});
    graph.instanceToIndex[instance] = task;
  }

  MergeScheduleResult pipelineSchedule;
  pipelineSchedule.processorCount = 2;
  pipelineSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
  for (size_t task = 0; task < graph.nodes.size(); ++task) {
    const ComputeInstance& instance = graph.nodes[task].instance;
    pipelineSchedule.dominanceOrderCompute.push_back(instance);
    size_t cpu = task < 4 ? 0 : 1;
    pipelineSchedule.computeToCpuMap[instance] = cpu;
    pipelineSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
    pipelineSchedule.computeToAestMap[instance] = task;
  }

  SchedulingTarget physical = fast;
  physical.processorCount = 4;
  physical.residentWeightCapacity = 2;
  physical.interProcessorLatencyNs = {
    0, 3, 3, 3,
    3, 0, 3, 3,
    3, 3, 0, 3,
    3, 3, 3, 0,
  };
  std::string pipelineError;
  assert(mlir::succeeded(applyPipelineScheduling(
    graph, pipelineSchedule, 2, physical, pipelineError)));
  assert(pipelineSchedule.processorCount == 4);
  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[0].instance) == 0);
  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[1].instance) == 0);
  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[2].instance) == 2);
  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[3].instance) == 2);
  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[4].instance) == 1);
  assert(pipelineSchedule.computeToCpuMap.lookup(graph.nodes[5].instance) == 3);
  assert(pipelineSchedule.computeToAestMap.lookup(graph.nodes[2].instance)
         >= pipelineSchedule.computeToAestMap.lookup(graph.nodes[1].instance)
              + graph.nodes[1].cost + 4);
  assert(pipelineSchedule.equivalentClass.empty());

  ComputeGraph communicationGraph;
  communicationGraph.successors.resize(5);
  communicationGraph.predecessors.resize(5);
  communicationGraph.successors[4].push_back(
    {3, TransferCost {.fixed = 0, .networkFlits = 1}});
  communicationGraph.predecessors[3].push_back(
    {4, TransferCost {.fixed = 0, .networkFlits = 1}});
  const Cost communicationCosts[] = {6, 4, 6, 4, 1};
  MergeScheduleResult communicationSchedule;
  communicationSchedule.processorCount = 2;
  for (uint32_t task = 0; task < 5; ++task) {
    ComputeInstance instance {nullptr, task, 1};
    ResidentWeight weight;
    weight.opaqueLane = task;
    communicationGraph.nodes.push_back(
      {instance, communicationCosts[task], {weight}, task});
    communicationGraph.instanceToIndex[instance] = task;
    communicationSchedule.dominanceOrderCompute.push_back(instance);
    size_t cpu = task < 4 ? 0 : 1;
    communicationSchedule.computeToCpuMap[instance] = cpu;
    communicationSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : 0;
    communicationSchedule.computeToAestMap[instance] = task;
  }

  SchedulingTarget fastPipeline = physical;
  fastPipeline.residentWeightCapacity = 4;
  MergeScheduleResult fastCommunicationSchedule = communicationSchedule;
  assert(mlir::succeeded(applyPipelineScheduling(
    communicationGraph, fastCommunicationSchedule, 2, fastPipeline, pipelineError)));
  assert(fastCommunicationSchedule.computeToCpuMap.lookup(
           communicationGraph.nodes[2].instance) == 2);

  SchedulingTarget slowPipeline = fastPipeline;
  slowPipeline.averageInterProcessorLatencyNs = 10;
  MergeScheduleResult slowCommunicationSchedule = communicationSchedule;
  assert(mlir::succeeded(applyPipelineScheduling(
    communicationGraph, slowCommunicationSchedule, 2, slowPipeline, pipelineError)));
  assert(slowCommunicationSchedule.computeToCpuMap.lookup(
           communicationGraph.nodes[2].instance) < 2);
  return EXIT_SUCCESS;
}