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

#define NUM_ROWS 8

using namespace tbb;
using namespace std;

typedef vector<vector<int>> Matrix;

struct DAGTask : public task
{
    vector<DAGTask*> successors;
    vector<int> receive;
    vector<int> data;
    vector<int*> gift;
    DAGTask* final_task; // samo je aktivan u zadnjij iteraciji ostalo nicemu ne sluzi
    int iteration; // po ovome razlikujem iteracije

    DAGTask(const vector<DAGTask*>& s, const vector<int>& rec, const vector<int*>& gifts, int it, const vector<int>& y, DAGTask* e = nullptr):
        successors(s),
        receive(NUM_ROWS),
        gift(gifts),
        final_task(e),
        iteration(it)
    {

        switch (iteration)
        {
        case 1:
            receive = rec;
            data = y;
            add_successors();
            break;
        case 2:
            data = rec;
            add_successors();
            break;
        case 3:
            data = rec;
            final_task->increment_ref_count();
            break;
        }
    }

    void add_successors() {
        for (DAGTask* dt : successors) {
            dt->increment_ref_count();
        }
    }

    task* execute()
    {
        int temp_sum = 0;
        switch (iteration)
        {
        case 1:
            for (size_t i = 0; i < NUM_ROWS; i++) { temp_sum += receive[i] * data[i]; }
            for (size_t i = 0; i < gift.size(); i++) { *(gift[i]) = temp_sum; }
            for (auto& a : successors)
            {
                if (a->decrement_ref_count() == 0) { spawn(*a); }
            }
            break;
        case 2:
            for (size_t i = 0; i < NUM_ROWS; i++) { temp_sum += receive[i] * data[i]; }
            for (size_t i = 0; i < gift.size(); i++) { *(gift[i]) = temp_sum; }
            for (auto& a : successors)
            {
                if (a->decrement_ref_count() == 0) { spawn(*a); }
            }
            break;
        case 3:
            for (size_t i = 0; i < NUM_ROWS; i++) { receive[i] *= data[i]; }
            for (size_t i = 0; i < NUM_ROWS; i++) { *(gift[i]) = receive[i]; }
            if (final_task->decrement_ref_count() == 0) {
                spawn(*final_task);
            }
            break;
        }
        return NULL;
    }
};


void BuildAndEvaluateDAG()
{
    // Inicijalizujem matrice
    vector<vector<int>> LN{ { 0,1,2,2,2,9,9,1 },{ 2,3,2,2,3,9,9,1 } ,{ 1,2,3,4,5,9,9,1 },{ 1,2,3,4,5,9,9,1 },{ 1,2,3,4,5,9,9,1 } ,{ 1,3,4,5,6,7,0,2 },{ 1,3,4,5,6,7,0,2 },{ 1,3,4,5,6,7,0,2 } };
    vector<vector<int>> C{ { 0,1,2,2,2,2,2,2 },{ 2,3,2,2,3,2,2,2 } ,{ 1,2,3,4,5,2,2,2 },{ 1,2,3,4,5,2,2,2 },{ 1,2,3,4,5,2,2,2 },{ 1,3,4,5,6,7,0,2 },{ 1,3,4,5,6,7,0,2 },{ 1,3,4,5,6,7,0,2 } };
    vector<vector<int>> ALPHA{ { 0,1,2,2,2,2,2,2 },{ 2,3,2,2,3,2,2,2 } ,{ 1,2,3,4,5,2,2,2 },{ 1,2,3,4,5,2,2,2 },{ 1,2,3,4,5 ,2,2,2 } ,{ 1,3,4,5,6,7,0,2 },{ 1,3,4,5,6,7,0,2 },{ 1,3,4,5,6,7,0,2 } };
    Matrix result;
    for (size_t i = 0; i < NUM_ROWS; ++i) { //formiram matricu za rezultat
        result.push_back(vector<int>(NUM_ROWS));
    }
    DAGTask* e = (DAGTask*) new(tbb::task::allocate_root()) empty_task();
    vector<DAGTask*> r_tasks;
    vector<DAGTask*> rr_tasks;
    vector<DAGTask*> rrr_tasks;

    for (size_t i = 0; i < NUM_ROWS; i++)
    {
        vector<int*> gifts_for_rrr;
        for (size_t j = 0; j < NUM_ROWS; j++)
        {
            gifts_for_rrr.push_back(&result[i][j]);
        }
        rrr_tasks.push_back(new (task::allocate_root()) DAGTask({}, ALPHA[i], move(gifts_for_rrr), 3, {}, e));

        for (size_t j = 0; j < NUM_ROWS; j++)
        {
            rr_tasks.push_back(new (task::allocate_root()) DAGTask({ rrr_tasks[i] }, C[j], { &(rrr_tasks[i]->receive[j]) }, 2, {}));
        }

        for (size_t j = 0; j < NUM_ROWS; j++)
        {
            vector<int*> gifts_for_rr;
            for (size_t n = i * NUM_ROWS; n < (i + 1) * NUM_ROWS; n++)
            {
                gifts_for_rr.push_back(&(rr_tasks[n]->receive[j]));
            }

            vector<int> whole_col;

            for (size_t i = 0; i < NUM_ROWS; i++)
            {
                whole_col.push_back(C[i][j]);
            }

            r_tasks.push_back(new (task::allocate_root()) DAGTask(vector<DAGTask*>(rr_tasks.begin() + i*NUM_ROWS, rr_tasks.begin() + (i + 1)*NUM_ROWS), LN[i], move(gifts_for_rr), 1, whole_col));
        }
    }

    e->increment_ref_count();
    for (DAGTask* r : r_tasks) {
        e->spawn(*r);
    }
    e->wait_for_all();
    e->execute();
    delete(e);

    for (size_t i = 0; i < NUM_ROWS; ++i)
    {
        for (size_t j = 0; j < NUM_ROWS; ++j) {
            cout << result[i][j] << " ";
        }
        cout << endl;
    }
}

int main(int argc, char* argv[])
{
    tick_count begin = tick_count::now();
    BuildAndEvaluateDAG();
    tick_count end = tick_count::now();
    cout << (end - begin).seconds()*1000 << "ms." << endl;
    return 0;
}