#pragma once

#include "../taskflow.hpp"

namespace tf {

class DeferredPipeflow {

};

class Pipeflow {

};

enum class PipeType : int {

};

template <typename C = std::function<void(tf::Pipeflow&)>>

class Pipe {

};

template <typename... Ps>

class Pipeline {

};

template <typename... Ps>

Pipeline<Ps...>::Pipeline(size_t num_lines, Ps&&... ps) :

_pipes {std::make_tuple(std::forward<Ps>(ps)...)},

_meta {PipeMeta{ps.type()}...},

_lines (num_lines),

_tasks (num_lines + 1),

_pipeflows (num_lines) {

if(num_lines == 0) {

TF_THROW("must have at least one line");

}

if(std::get<0>(_pipes).type() != PipeType::SERIAL) {

TF_THROW("first pipe must be serial");

}

reset();

_build();

}

template <typename... Ps>

Pipeline<Ps...>::Pipeline(size_t num_lines, std::tuple<Ps...>&& ps) :

_pipes {std::forward<std::tuple<Ps...>>(ps)},

_meta {_gen_meta(

std::forward<std::tuple<Ps...>>(ps), std::make_index_sequence<sizeof...(Ps)>{}

)},

_lines (num_lines),

_tasks (num_lines + 1),

_pipeflows (num_lines) {

if(num_lines == 0) {

TF_THROW("must have at least one line");

}

if(std::get<0>(_pipes).type() != PipeType::SERIAL) {

TF_THROW("first pipe must be serial");

}

reset();

_build();

}

template <typename... Ps>

template <size_t... I>

auto Pipeline<Ps...>::_gen_meta(std::tuple<Ps...>&& ps, std::index_sequence<I...>) {

return std::array{PipeMeta{std::get<I>(ps).type()}...};

}

template <typename... Ps>

size_t Pipeline<Ps...>::num_lines() const noexcept {

return _pipeflows.size();

}

template <typename... Ps>

constexpr size_t Pipeline<Ps...>::num_pipes() const noexcept {

return sizeof...(Ps);

}

template <typename... Ps>

size_t Pipeline<Ps...>::num_tokens() const noexcept {

return _num_tokens;

}

template <typename... Ps>

Graph& Pipeline<Ps...>::graph() {

return _graph;

}

template <typename... Ps>

void Pipeline<Ps...>::reset() {

_num_tokens = 0;

for(size_t l = 0; l<num_lines(); l++) {

_pipeflows[l]._pipe = 0;

_pipeflows[l]._line = l;

_pipeflows[l]._num_deferrals = 0;

_pipeflows[l]._dependents.clear();

}

assert(_ready_tokens.empty() == true);

_token_dependencies.clear();

_deferred_tokens.clear();

_lines[0][0].join_counter.store(0, std::memory_order_relaxed);

for(size_t l=1; l<num_lines(); l++) {

for(size_t f=1; f<num_pipes(); f++) {

_lines[l][f].join_counter.store(

static_cast<size_t>(_meta[f].type), std::memory_order_relaxed

);

}

}

for(size_t f=1; f<num_pipes(); f++) {

_lines[0][f].join_counter.store(1, std::memory_order_relaxed);

}

for(size_t l=1; l<num_lines(); l++) {

_lines[l][0].join_counter.store(

static_cast<size_t>(_meta[0].type) - 1, std::memory_order_relaxed

);

}

}

template <typename... Ps>

void Pipeline<Ps...>::_on_pipe(Pipeflow& pf, Runtime& rt) {

visit_tuple([&](auto&& pipe){

using callable_t = typename std::decay_t<decltype(pipe)>::callable_t;

if constexpr (std::is_invocable_v<callable_t, Pipeflow&>) {

pipe._callable(pf);

}

else if constexpr(std::is_invocable_v<callable_t, Pipeflow&, Runtime&>) {

pipe._callable(pf, rt);

}

else {

static_assert(dependent_false_v<callable_t>, "un-supported pipe callable type");

}

}, _pipes, pf._pipe);

}

template <typename... Ps>

void Pipeline<Ps...>::_check_dependents(Pipeflow& pf) {

//if (pf._dependents.size()) {

++pf._num_deferrals;

for (auto it = pf._dependents.begin(); it != pf._dependents.end();) {

// valid (e.g., 12.defer(16))

if (*it >= _num_tokens) {

_token_dependencies[*it].push_back(pf._token);

_longest_deferral = std::max(_longest_deferral, *it);

++it;

}

// valid or invalid (e.g., 12.defer(7))

else {

auto pit = _deferred_tokens.find(*it);

// valid (e.g., 7 is deferred)

if (pit != _deferred_tokens.end()) {

_token_dependencies[*it].push_back(pf._token);

++it;

}

// invalid (e.g., 7 is finished - this this 12.defer(7) is dummy)

else {

it = pf._dependents.erase(it);

}

}

}

}

template <typename... Ps>

void Pipeline<Ps...>::_construct_deferred_tokens(Pipeflow& pf) {

//auto res = _deferred_tokens.emplace(

// pf._token, DeferredPipeflow{pf._token, pf._num_deferrals, std::move(pf._dependents)}

//);

// construct the deferred pipeflow with zero copy

//auto res = _deferred_tokens.emplace(

_deferred_tokens.emplace(

std::piecewise_construct,

std::forward_as_tuple(pf._token),

std::forward_as_tuple(

pf._token, pf._num_deferrals, std::move(pf._dependents)

)

);

//assert(res.second == true);

}

template <typename... Ps>

void Pipeline<Ps...>::_resolve_token_dependencies(Pipeflow& pf) {

if (auto it = _token_dependencies.find(pf._token);

it != _token_dependencies.end()) {

// iterate tokens that defer to pf._token

// (e.g., 12 and 13)

for(size_t target : it->second) {

auto dpf = _deferred_tokens.find(target);

assert(dpf != _deferred_tokens.end());

// erase pf._token from target's _dependents

// (e.g., remove 16 from 12's dependents)

dpf->second._dependents.erase(pf._token);

// dpf->second._dependent_satellites[pf._token]

//);

// target has no dependents

if (dpf->second._dependents.empty()) {

// push target into _ready_tokens queue

_ready_tokens.emplace(dpf->second._token, dpf->second._num_deferrals);

//_ready_tokens.push(

// std::make_pair(dpf->second._token, dpf->second._num_deferrals)

//);

// erase target from _deferred_tokens

_deferred_tokens.erase(dpf);

}

}

// remove pf._token from _token_dependencies

// (e.g., remove the entry

// {key: 16, value: std::vector{12, 13}} from _token_dependencies)

_token_dependencies.erase(it);

}

}

template <typename... Ps>

void Pipeline<Ps...>::_build() {

using namespace std::literals::string_literals;

FlowBuilder fb(_graph);

// init task

_tasks[0] = fb.emplace([this]() {

return static_cast<int>(_num_tokens % num_lines());

}).name("cond");

// line task

for(size_t l = 0; l < num_lines(); l++) {

_tasks[l + 1] = fb.emplace([this, l] (tf::Runtime& rt) mutable {

auto pf = &_pipeflows[l];

pipeline:

_lines[pf->_line][pf->_pipe].join_counter.store(

static_cast<size_t>(_meta[pf->_pipe].type), std::memory_order_relaxed

);

// First pipe does all jobs of initialization and token dependencies

if (pf->_pipe == 0) {

// _ready_tokens queue is not empty

// substitute pf with the token at the front of the queue

if (!_ready_tokens.empty()) {

pf->_token = _ready_tokens.front().first;

pf->_num_deferrals = _ready_tokens.front().second;

_ready_tokens.pop();

}

else {

pf->_token = _num_tokens;

pf->_num_deferrals = 0;

}

handle_token_dependency:

if (pf->_stop = false, _on_pipe(*pf, rt); pf->_stop == true) {

// here, the pipeline is not stopped yet because other

// lines of tasks may still be running their last stages

return;

}

if (_num_tokens == pf->_token) {

++_num_tokens;

}

if (pf->_dependents.empty() == false){

// check if the pf->_depdendents have valid dependents

_check_dependents(*pf);

// tokens in pf->_dependents are all valid dependents

if (pf->_dependents.size()) {

// construct a data structure for pf in _deferred_tokens

_construct_deferred_tokens(*pf);

goto pipeline;

}

// tokens in pf->_dependents are invalid dependents

// directly goto on_pipe on the same line

else {

goto handle_token_dependency;

}

}

// Every token within the deferral range needs to check

// if it can resolve dependencies on other tokens.

if (pf->_token <= _longest_deferral) {

_resolve_token_dependencies(*pf);

}

}

else {

_on_pipe(*pf, rt);

}

size_t c_f = pf->_pipe;

size_t n_f = (pf->_pipe + 1) % num_pipes();

size_t n_l = (pf->_line + 1) % num_lines();

pf->_pipe = n_f;

// ---- scheduling starts here ----

// Notice that the shared variable f must not be changed after this

// point because it can result in data race due to the following

// condition:

//

// a -> b

// | |

// v v

// c -> d

//

// d will be spawned by either c or b, so if c changes f but b spawns d

// then data race on f will happen

std::array<int, 2> retval;

size_t n = 0;

// downward dependency

if(_meta[c_f].type == PipeType::SERIAL &&

_lines[n_l][c_f].join_counter.fetch_sub(

1, std::memory_order_acq_rel) == 1

) {

retval[n++] = 1;

}

// forward dependency

if(_lines[pf->_line][n_f].join_counter.fetch_sub(

1, std::memory_order_acq_rel) == 1

) {

retval[n++] = 0;

}

// notice that the task index starts from 1

switch(n) {

case 2: {

rt.schedule(_tasks[n_l+1]);

goto pipeline;

}

case 1: {

// downward dependency

if (retval[0] == 1) {

pf = &_pipeflows[n_l];

}

// forward dependency

goto pipeline;

}

}

}).name("rt-"s + std::to_string(l));

_tasks[0].precede(_tasks[l+1]);

}

}