#include <iostream>
#include <stdlib.h>
#include <string>
#include "tbb/task_scheduler_init.h"
#include "tbb/tick_count.h"
#include "tbb/task.h"
#include <vector>

using namespace tbb;
using namespace std;

const int M = 2, N = 2;

class DagTask : public tbb::task {
public:

    // input[0] = sum from above, input[1] = sum from left
    vector<int> input;
    vector<int> data;
    string id;
    int type;
    // successor[0] = successor below, successor[1] = successor to right
    vector<DagTask*> successor;
    DagTask(vector<int> _data, vector<int> _input, string _id, int t = 0) :  data(_data), input(_input), id(_id), type(t){}

    void add_successor(DagTask* s) {
        successor.push_back(s);
        s->increment_ref_count();
    }

    task* execute() {
        __TBB_ASSERT(ref_count() == 0, NULL);
        printf_s("%s %d %d %d %d\n", id.data(), data[0], data[1], input[0], input[1]);

        if (type == 1) {
            /*printf_s("%s se sjebo\n", id.data());*/
            for (size_t i = 0; i < input.size(); i++)
            {
                data[i] *= input[input.size()-1-i];
                successor[0]->input.push_back(data[i]);
            }
            printf_s("%s %d %d %d %d\n", id.data(), data[0], data[1], input[0], input[1]);
        }

        for (int k = 0; k < successor.size(); ++k)
            if (DagTask* t = successor[k]) {
                t->input.push_back(data[0] * input[1] + data[1] * input[0]);
                if (t->decrement_ref_count() == 0)
                    spawn(*t);
            }

        return NULL;
    }
};

void BuildAndEvaluateDAG() {

    DagTask* e = (DagTask*) new(tbb::task::allocate_root()) empty_task;

    DagTask* rrr1 = new(tbb::task::allocate_root()) DagTask({ 8, 9 }, {}, "rrr1", 1);
    rrr1->add_successor(e);
    DagTask* rrr2 = new(tbb::task::allocate_root()) DagTask({ 10, 11 }, {}, "rrr2", 1);
    rrr2->add_successor(e);

    DagTask* rr1 = new(tbb::task::allocate_root()) DagTask({ 4, 5 }, {}, "rr1");
    rr1->add_successor(rrr1);
    DagTask* rr2 = new(tbb::task::allocate_root()) DagTask({ 6, 7 }, {}, "rr2");
    rr2->add_successor(rrr1);
    DagTask* rr3 = new(tbb::task::allocate_root()) DagTask({ 4, 5 }, {}, "rr3");
    rr3->add_successor(rrr2);
    DagTask* rr4 = new(tbb::task::allocate_root()) DagTask({ 6, 7 }, {}, "rr4");
    rr4->add_successor(rrr2);

    DagTask* r1 = new(tbb::task::allocate_root()) DagTask({ 0, 1 },  { 6, 4 }, "r1");
    r1->add_successor(rr1);
    r1->add_successor(rr2);
    DagTask* r2 = new(tbb::task::allocate_root()) DagTask({ 0, 1 }, { 7, 5 }, "r2");
    r2->add_successor(rr1);
    r2->add_successor(rr2);
    DagTask* r3 = new(tbb::task::allocate_root()) DagTask({ 2, 3 }, { 6, 4 }, "r3");
    r3->add_successor(rr3);
    r3->add_successor(rr4);
    DagTask* r4 = new(tbb::task::allocate_root()) DagTask({ 2, 3 }, { 7, 5 }, "r4");
    r4->add_successor(rr3);
    r4->add_successor(rr4);


    e->increment_ref_count();
    e->spawn(*r1);
    e->spawn(*r2);
    e->spawn(*r3);
    e->spawn_and_wait_for_all(*r4);
    cout << "zavrsio" << endl;
    e->execute();
    int result[2][2];
    result[0][0] = e->input[3];
    result[0][1] = e->input[2];
    result[1][0] = e->input[0];
    result[1][1] = e->input[1];
    for (size_t i = 0; i < 2; i++)
    {
        for (size_t j = 0; j < 2; j++)
        {
            cout << result[i][j] << " ";
        }
        cout << endl;
    }
    task::destroy(*e);

}

int main(int argc, char* argv[])
{
    // build and evaluate DAG
    BuildAndEvaluateDAG();
    return 0;
}