
#ifndef PLS_DATAFLOW_GRAPH_H_
#define PLS_DATAFLOW_GRAPH_H_

#include "inputs.h"
#include "outputs.h"
#include "input.h"
#include "output.h"

#include "internal/token.h"

namespace pls {
namespace dataflow {

template<typename I, typename O, int P>
class graph {
  /**
   * Boilerplate for callbacks of dataflow input/output handlers.
   */
  struct external_input_cb {
    graph *graph_;

    void one_input(int pos, internal::token_color color) {
      graph_->one_external_input(pos, color);
    }
    void all_inputs(internal::token_color color) {
      graph_->all_external_inputs(color);
    }
  };
  struct internal_output_cb {
    graph *graph_;

    void one_input(int pos, internal::token_color color) {
      graph_->one_internal_output(pos, color);
    }
    void all_inputs(internal::token_color color) {
      graph_->all_internal_outputs(color);
    }
  };

  /**
   * Handling of buffers for dataflow.
   */
  using exeternal_inputs = typename I::template internal_inputs<P, graph<I, O, P>::external_input_cb>;
  using external_outputs = typename O::template internal_outputs<P>;
  exeternal_inputs external_inputs_{external_input_cb{this}};
  external_outputs external_outputs_{};

  void one_external_input(int pos, internal::token_color color) {
    // We do currently not support 'single external inputs'.
    // This prevents some fine-grained parallelism, but eases implementation
  }
  void all_external_inputs(internal::token_color external_color) {
    auto my_clock = input_clock_++;
    internal::token_color internal_color{my_clock, external_color.depth_ + 1};
    auto my_index = internal_color.get_index(P);

    while (output_clear_[my_index] != clear) {
      PLS_ERROR("Full inputs are not implemented jet. Place waiting code in here/yield to scheduler.")
    }

    output_clear_[my_index] = waiting;
    input_output_color_matching_[my_index] = {external_color, internal_color};
    external_inputs_.push_to_outputs(internal_inputs_, external_color.get_index(P), internal_color);
  }

  using internal_outputs = typename I::template internal_inputs<P, graph<I, O, P>::internal_output_cb>;
  using internal_inputs = typename O::template internal_outputs<P>;
  internal_outputs internal_outputs_{internal_output_cb{this}};
  internal_inputs internal_inputs_{};

  void one_internal_output(int pos, internal::token_color color) {
    // We do not support 'single flow' output. Every graph is well behaved, e.g. will emmit exactly one output tuple.
  }
  void all_internal_outputs(internal::token_color color) {
    auto color_matching = input_output_color_matching_[color.get_index(P)];
    auto external_color = std::get<0>(color_matching);
    auto internal_color = std::get<1>(color_matching);
    if (external_color == internal_color) {
      output_clear_[color.get_index(P)] = finished;
    } else {
      internal_outputs_.push_to_outputs(external_outputs_, internal_color.get_index(P), external_color);
      output_clear_[color.get_index(P)] = clear;
    }
  }

  /**
   * Handling of colors/maximum concurrent tokens in the system.
   */
  std::atomic<unsigned int> input_clock_{0};
  std::atomic<unsigned int> output_clock_{0};
  std::array<std::pair<internal::token_color, internal::token_color>, P> input_output_color_matching_{};
  enum output_state { clear, waiting, finished };
  std::array<output_state, P> output_clear_{clear};

 public:
  template<int N>
  using external_output_at = output<P, typename external_outputs::template raw_type_at<N>>;
  template<int N>
  using external_input_at = input<P, typename exeternal_inputs::template raw_type_at<N>>;

  template<int N>
  external_output_at<N> external_output() {
    external_output_at<N> o{external_outputs_.template get<N>()};
    return o;
  }
  template<int N>
  external_input_at<N> external_input() {
    external_input_at<N> i{external_inputs_.template get<N>()};
    return i;
  }

  template<int N>
  using input_at = output<P, typename internal_inputs::template raw_type_at<N>>;
  template<int N>
  using output_at = input<P, typename internal_outputs::template raw_type_at<N>>;

  template<int N>
  output_at<N> output() {
    output_at<N> o{internal_outputs_.template get<N>()};
    return o;
  }
  template<int N>
  input_at<N> input() {
    input_at<N> i{internal_inputs_.template get<N>()};
    return i;
  }

  template<typename ...ARGS>
  void push_input(ARGS ...values) {
    auto my_clock = input_clock_++;
    internal::token_color color{my_clock, 0};
    auto my_index = color.get_index(P);

    while (output_clear_[my_index] != clear) {
      PLS_ERROR("Full inputs are not implemented jet. Place waiting code in here/yield to scheduler.")
    }

    output_clear_[my_index] = waiting;
    input_output_color_matching_[my_index] = {color, color};

    push_input_rec<0, ARGS...>(color, values...);
  }
  template<int POS, typename ...TAIL>
  void push_input_rec(internal::token_color color) {}
  template<int POS, typename HEAD, typename ...TAIL>
  void push_input_rec(internal::token_color color, HEAD h, TAIL ...t) {
    internal_inputs_.template get<POS>().push_token(internal::token<HEAD>{h, color});
    push_input_rec<POS + 1, TAIL...>(color, t...);
  }

  typename internal_outputs::raw_types get_output() {
    while (true) {
      if (output_clock_ > input_clock_) {
        PLS_ERROR("Must not try to pull outputs if there are not enough inputs.")
      }

      auto color_matching = input_output_color_matching_[output_clock_ % P];
      if (std::get<0>(color_matching) == std::get<1>(color_matching)) {
        break;
      }
      output_clock_++;
    }

    while (output_clear_[output_clock_ % P] != finished) {
      PLS_ERROR("Blocking/Unfinishe outputs are not implemented jet. Place waiting code in here/yield to scheduler.")
    }

    auto result = internal_outputs_.get_outputs(output_clock_ % P);
    output_clear_[output_clock_ % P] = clear;
    output_clock_++;
    return result;
  }
};

}
}

#endif //PLS_DATAFLOW_GRAPH_H_
