#include "../context.h"
#include "Iterator.h"
#include "TaskManager.h"
#include <stdio.h>

Iterator* createMultiDimIterator(struct Context* context, int x, int y, int z) {
    struct Iterator* iterator = new Iterator();
    struct MultiDimIterator* multiDimIterator = new MultiDimIterator();
    iterator->iterator = (union Data*)multiDimIterator;
    iterator->exec = C_execMultiDimIterator;
    iterator->barrier = C_barrierMultiDimIterator;
    multiDimIterator->x = x;
    multiDimIterator->y = y;
    multiDimIterator->z = z;
    multiDimIterator->count = x * y * z;
    multiDimIterator->counterX = 0;
    multiDimIterator->counterY = 0;
    multiDimIterator->counterZ = 0;
    return iterator;
}

/**
 * create iterateTask with index, that copy from task argument
 * @return created iterateTask
 * @param task      task of the copy source
 * @x     index
 */
struct Context* createMultiDimIterateTask(struct Context* task, int x, int y, int z) {
    struct Context* task1 = NEW(struct Context);
    initContext(task1);
    task1->taskManager = task->taskManager;
    task1->next     = task->next;
    task1->iterate  = 1;
    task1->iterator = task->iterator;
    task1->idgCount = task->idgCount;
    task1->idg      = task->idg;
    task1->maxIdg   = task->maxIdg;
    for(int i = task1->idg; i < task1->maxIdg; i++) {
        task1->data[i] = task->data[i];
    }

    // create index data gear and register input data to iterate task
    struct MultiDim* multiDim = &ALLOCATE_DATA_GEAR(task1, MultiDim)->MultiDim;
    multiDim->x = x;
    multiDim->y = y;
    multiDim->z = z;
    task1->data[task1->maxIdg++] = (union Data*)multiDim;
    task1->odg      = task->odg + 1;
    task1->maxOdg   = task->maxOdg + 1;
    for (int i = task1->odg; i < task1->maxOdg; i++) {
        task1->data[i] = task->data[i-1];
    }

    return task1;
}

__code execMultiDimIterator(struct MultiDimIterator* iterator, struct Context* task, int numGPU, __code next(...)) {
    // No GPU device
    if (numGPU == 0) {
        goto meta(context, C_execMultiDimIterator1);
    }
    task->iterate = 1;
    task->gpu = 1;
    struct TaskManager* taskManager = task->taskManager;
    goto taskManager->spawn(task, next(...));
}

__code execMultiDimIterator_stub(struct Context* context) {
    MultiDimIterator* iterator = (MultiDimIterator*)GearImpl(context, Iterator, iterator);
    Context* task = Gearef(context, Iterator)->task;
    int numGPU = Gearef(context, Iterator)->numGPU;
    enum Code next = Gearef(context, Iterator)->next;
    goto execMultiDimIterator(context, iterator, task, numGPU, next);
} 

__code execMultiDimIterator1(struct MultiDimIterator* iterator, struct Context* task, __code next(...)) {
    int x = iterator->counterX;
    int y = iterator->counterY;
    int z = iterator->counterZ;
    struct Context* iterateTask = createMultiDimIterateTask(task, x, y, z);
    struct TaskManager* taskManager = task->taskManager;
    goto taskManager->spawn(iterateTask, execMultiDimIterator2);
}

__code execMultiDimIterator1_stub(struct Context* context) {
    MultiDimIterator* iterator = (MultiDimIterator*)GearImpl(context, Iterator, iterator);
    Context* task = Gearef(context, Iterator)->task;
    enum Code next = Gearef(context, Iterator)->next;
    goto execMultiDimIterator1(context, iterator, task, next);
} 

__code execMultiDimIterator2(struct MultiDimIterator* iterator, struct Context* task, __code next(...)) {
    if (++iterator->counterX >= iterator->x) {
        iterator->counterX = 0;
        if (++iterator->counterY >= iterator->y) {
            iterator->counterY = 0;
            if (++iterator->counterZ >= iterator->z) {
                iterator->counterZ = 0;
                goto next(...);
            }
        }
    }
    goto meta(context, C_execMultiDimIterator1);
}

__code barrierMultiDimIterator(struct MultiDimIterator* iterator, struct Context* task, __code next(...), __code whenWait(...)) {
    if (task->gpu || __sync_fetch_and_sub(&iterator->count, 1) == 1) {
        goto next(...);
    }
    goto whenWait(...);
}

__code barrierMultiDimIterator_stub(struct Context* context) {
    MultiDimIterator* iterator = (MultiDimIterator*)GearImpl(context, Iterator, iterator);
    Context* task = Gearef(context, Iterator)->task;
    enum Code next = Gearef(context, Iterator)->next;
    enum Code whenWait = Gearef(context, Iterator)->whenWait;
    goto barrierMultiDimIterator(context, iterator, task, next, whenWait);
}