// 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/switch_node.h"
#include "pls/dataflow/internal/split_node.h"
#include "pls/dataflow/internal/merge_node.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_1{minus_one};
  function_node<inputs<int>, outputs<int>, decltype(minus_one)> minus_one_node_2{minus_one};

  auto is_positive = [](const int &i1, bool &o1) {
    o1 = i1 > 0;
  };
  function_node<inputs<int>, outputs<bool>, decltype(is_positive)> is_positive_node{is_positive};

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

  split_node<int> minus_split;
  split_node<bool> decision_split;
  switch_node<int> recursion_split;
  merge_node<int> recursion_merge;

  // Connect
  minus_one_node_1 >> minus_one_node_2;

  // Inputs to first processing step
  graph.input<0>() >> minus_one_node_1.in_port<0>();
  graph.input<1>() >> triple_node.in_port<0>();
  minus_one_node_2.out_port<0>() >> minus_split.value_in_port();

  // Prepare decision...
  minus_split.out_port_1() >> is_positive_node.in_port<0>();
  is_positive_node.out_port<0>() >> decision_split.value_in_port();
  triple_node.out_port<0>() >> recursion_split.value_in_port();
  decision_split.out_port_1() >> recursion_split.condition_in_port();

  // Connect true case (recursion)
  minus_split.out_port_2() >> recursion_node.in_port<0>();
  recursion_split.true_out_port() >> recursion_node.in_port<1>();
  recursion_node.out_port<0>() >> recursion_merge.true_in_port();

  // Connect false case (no recursion)
  recursion_split.false_out_port() >> recursion_merge.false_in_port();

  // Deliver final result (back from merge)
  decision_split.out_port_2() >> recursion_merge.condition_in_port();
  recursion_merge.value_out_port() >> 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;
  });
}
