// Headers are available because we added the pls target
#include <string>
#include <cstdio>
#include <tuple>
#include <array>

#include <pls/pls.h>
#include <pls/dataflow/internal/inputs.h>
#include <pls/dataflow/internal/outputs.h>
#include <pls/dataflow/internal/function_node.h>
#include <pls/dataflow/internal/graph.h>
#include <pls/dataflow/internal/out_port.h>

int main() {
  using namespace pls::dataflow;
  using namespace pls::dataflow::internal;

  // Define
  graph<inputs<int, int>, outputs<int>> graph;

  auto triple = [](const int &i1, int &o1) {
    o1 = i1 * 3;
  };
  function_node<inputs<int>, outputs<int>, decltype(triple)> triple_node{triple};

  auto minus_one = [](const int &i1, int &o1) {
    o1 = i1 - 1;
  };
  function_node<inputs<int>, outputs<int>, decltype(minus_one)> minus_one_node{minus_one};

  auto recursion = [&](const int &i1, const int &i2, int &o1) {
    if (i1 > 0) {
      std::tuple<int> out;
      graph.run({i1, i2}, out);
      pls::scheduler::wait_for_all();
      o1 = std::get<0>(out);
    } else {
      o1 = i2;
    }
  };
  function_node<inputs<int, int>, outputs<int>, decltype(recursion)> recursion_node{recursion};

  // Connect
  graph.input<0>() >> minus_one_node.in_port<0>();
  minus_one_node.out_port<0>() >> recursion_node.in_port<0>();

  graph.input<1>() >> triple_node.in_port<0>();
  triple_node.out_port<0>() >> recursion_node.in_port<1>();

  recursion_node.out_port<0>() >> graph.output<0>();

  // Build
  graph.build();

  pls::malloc_scheduler_memory my_scheduler_memory{8, 2u << 18u};
  pls::scheduler scheduler{&my_scheduler_memory, 8};
  scheduler.perform_work([&] {
    // Schedule Execution
    std::tuple<int> out1, out2, out3;
    graph.run({1, 2}, out1);
    graph.run({1, 1}, out2);
    graph.run({5, 6}, out3);

    // Wait for results and print
    pls::scheduler::wait_for_all();
    std::cout << std::get<0>(out1) << std::endl;
    std::cout << std::get<0>(out2) << std::endl;
    std::cout << std::get<0>(out3) << std::endl;
  });
}
