#pragma once

#include "llvm/ADT/ArrayRef.h"

#include <list>
#include <optional>
#include <unordered_map>
#include <vector>

#include "DCPAnalysis.hpp"
#include "Task.hpp"
#include "Utils.hpp"

std::optional<Edge_pair> addEdge(TaskDCP* parent, TaskDCP* child, Weight_t weight);
void removeEdge(TaskDCP* parent, TaskDCP* child);
int getTranferCost(TaskDCP* parent, TaskDCP* child);

class GraphDCP {

  struct FindSlot {
    int aest;
    int index;
  };

  std::vector<TaskDCP> nodes;
  std::unordered_map<CPU, std::list<TaskDCP*>> mapCPUTasks;
  CPU last_cpu = 0;

  std::array<std::optional<Edge_pair>, 2> insertTaskInCPU(CPU cpu, TaskDCP* task, size_t position);
  void removeTaskFromCPU(CPU cpu, TaskDCP* task);

  std::vector<TaskDCP*> getRoots();

  void initAEST();
  int computeAEST(TaskDCP* task, CPU cpu);
  int initDCPL();
  int computeDCPL(TaskDCP* task, CPU cpu);
  void initALST();
  std::unordered_map<TaskDCP*, int> computeALST(TaskDCP* task, CPU cpu);

  TaskDCP* findCandidate(std::vector<TaskDCP*> nodes);
  void selectProcessor(TaskDCP* candidate, bool push);
  CPU lastCPU() const { return last_cpu; }
  void incLastCPU() { last_cpu++; }
  FindSlot findSlot(TaskDCP* candidate, CPU cpu, bool push);
  void to_dot();

public:
  void DCP();
  GraphDCP(llvm::ArrayRef<onnx_mlir::spatial::SpatWeightedCompute> spatWeightedComputes,
           llvm::ArrayRef<EdgesIndex> edges)
  : nodes(), mapCPUTasks() {
    for (auto spatWeightedCompute : spatWeightedComputes)
      nodes.emplace_back(spatWeightedCompute);

    for (auto [start, end, weight] : edges)
      makeEdge(start, end, weight);
  }

  DCPAnalysisResult getResult();

  void makeEdge(size_t parent_index, size_t child_index, Weight_t weight) {
    addEdge(&nodes[parent_index], &nodes[child_index], weight);
  }

  size_t taskInCPU(CPU cpu) { return mapCPUTasks[cpu].size(); }
};