#include "def.h"
#include "simple.h"
#include "cbag.h"

// necessary external definition to tell OpenMP that head is shared state
extern simple_node_t *head;
// global/shared cbag state
extern block_t **global_head_block;

#if VERIFY
int *start_indices;
int *items_to_add;
bool *processed_correctly;
#endif

/* These structs should to match the definition in benchmark.py
 */
struct counters
{
    int additions;
    int failed_removals;
    int successful_removals;
};
struct bench_result
{
    float time;
    float add_time;
    float try_remove_any_time;
    struct counters reduced_counters;
};

void noop() {}
void noop_with_int(__attribute__((unused)) int unused) {}

struct counters bench_simple_prod_or_cons_thread(int workload, bool is_consumer)
{
    int tid = omp_get_thread_num();
    // printf("thread %2d started\n", tid);

// barrier to force OMP to start all threads at the same time
#pragma omp barrier

    struct counters data = {.additions = 0, .failed_removals = 0, .successful_removals = 0};

    for (int i = 0; i < workload; i++)
    {
#pragma omp critical
        {
            if (is_consumer)
            {
                void *item = simple_try_remove_any();
                if (item != NULL)
                {
#if VERIFY
                    int result = *((int *)item);
                    processed_correctly[result] = true;
#else
                    free(item);
#endif
                    data.successful_removals++;
                }
                else
                {
                    data.failed_removals++;
                    i--;
                }
            }
            else
            {
                int *item;
#if VERIFY
                item = &(items_to_add[start_indices[tid] + i]);
#else
                item = malloc(sizeof(int));
                // prevent NULL from being written if t0 is a producer
                *item = tid + 1;
#endif
                simple_add((void *)item);
                data.additions++;
            }
        }
    }

    return data;
}

struct counters bench_cbag_prod_or_cons_thread(int workload, bool is_consumer)
{
    int tid = omp_get_thread_num();
    // printf("thread %2d started\n", tid);

    cbag_init_thread();

// barrier to force OMP to start all threads at the same time
#pragma omp barrier

    struct counters data = {.additions = 0, .failed_removals = 0, .successful_removals = 0};

    for (int i = 0; i < workload; i++)
    {
        if (is_consumer)
        {
            void *item = cbag_try_remove_any();
            if (item != NULL)
            {
#if VERIFY
                int result = *((int *)item);
#pragma omp critical
                processed_correctly[result] = true;
#else
                free(item);
#endif
                data.successful_removals++;
            }
            else
            {
                data.failed_removals++;
                i--;
            }
        }
        else
        {
            int *item;
#if VERIFY
            item = &(items_to_add[start_indices[tid] + i]);
#else
            item = malloc(sizeof(int));
            // prevent NULL from being written if t0 is a producer
            *item = tid + 1;
#endif
            cbag_add((void *)item);
            data.additions++;
        }
    }

    return data;
}

struct counters bench_simple_add_thread(int workload)
{
    int tid = omp_get_thread_num();
    // printf("thread %2d started\n", tid);

// barrier to force OMP to start all threads at the same time
#pragma omp barrier

    struct counters data = {.additions = 0, .failed_removals = 0, .successful_removals = 0};

    for (int i = 0; i < workload; i++)
    {
#pragma omp critical
        {
            int *item;
#if VERIFY
            item = &(items_to_add[start_indices[tid] + i]);
#else
            item = malloc(sizeof(int));
            // prevent NULL from being written if t0 is a producer
            *item = tid + 1;
#endif
            simple_add((void *)item);
            data.additions++;
        }
    }

    return data;
}

struct counters bench_cbag_add_thread(int workload)
{
    int tid = omp_get_thread_num();
    // printf("thread %2d started\n", tid);

    cbag_init_thread();

// barrier to force OMP to start all threads at the same time
#pragma omp barrier

    struct counters data = {.additions = 0, .failed_removals = 0, .successful_removals = 0};

    for (int i = 0; i < workload; i++)
    {
        int *item;
#if VERIFY
        item = &(items_to_add[start_indices[tid] + i]);
#else
        item = malloc(sizeof(int));
        // prevent NULL from being written if t0 is a producer
        *item = tid + 1;
#endif
        cbag_add((void *)item);
        data.additions++;
    }

    return data;
}

struct counters bench_simple_try_remove_any_thread(int workload)
{
    // int tid = omp_get_thread_num();
    //  printf("thread %2d started\n", tid);

// barrier to force OMP to start all threads at the same time
#pragma omp barrier

    struct counters data = {.additions = 0, .failed_removals = 0, .successful_removals = 0};

    for (int i = 0; i < workload; i++)
    {
#pragma omp critical
        {
            void *item = simple_try_remove_any();
            if (item != NULL)
            {
#if VERIFY
                int result = *((int *)item);
                processed_correctly[result] = true;
#else
                free(item);
#endif
                data.successful_removals++;
            }
            else
            {
                data.failed_removals++;
                i--;
            }
        }
    }

    return data;
}

struct counters bench_cbag_try_remove_any_thread(int workload)
{
    // int tid = omp_get_thread_num();
    //  printf("thread %2d started\n", tid);

    cbag_init_thread();

// barrier to force OMP to start all threads at the same time
#pragma omp barrier

    struct counters data = {.additions = 0, .failed_removals = 0, .successful_removals = 0};

    for (int i = 0; i < workload; i++)
    {
        void *item = cbag_try_remove_any();
        if (item != NULL)
        {
#if VERIFY
            int result = *((int *)item);
#pragma omp critical
            processed_correctly[result] = true;
#else
            free(item);
#endif
            data.successful_removals++;
        }
        else
        {
            data.failed_removals++;
            i--;
        }
    }

    return data;
}

struct bench_result bench_prod_cons(int t, int items, int cons, char *name, void (*init_function)(int), void (*destroy_function)(), struct counters (*bench_prod_or_cons_thread_function)(int, bool))
{
    struct bench_result result;
    struct counters thread_data[t];
    double tic, toc;
    int prod = t - cons;

    if (t < 1 || items < 1 || cons < 1)
    {
        printf("invalid bench_prod_cons call: thread count, problem size and consumer count must be greater than 1\n");
        printf("    thread count = %d, problem size = %d, consumer count = %d, producer count = %d\n", t, items, cons, prod);
        return result;
    }

    printf(
        "benchmarking %s producer/consumer with %d consumers and %d producers out of %d threads "
        "with problem size %d\n",
        name, cons, prod, t, items);

    // first producer as well as first consumer need to handle excess elements
    int t0_cons_workload = items / cons + items % cons;
    int t0_prod_workload = items / prod + items % prod;
    // all the other take the neatly divisible parts of the work
    int t_cons_workload = items / cons;
    int t_prod_workload = items / prod;

#if VERIFY
    start_indices = (int *)malloc(sizeof(int) * t);
    items_to_add = (int *)malloc(sizeof(int) * items);
    processed_correctly = (bool *)malloc(sizeof(bool) * items);
    for (int i = 0; i < items; i++)
    {
        items_to_add[i] = i;
        processed_correctly[i] = false;
    }
    // 0th thread starts working (producing) at index 0
    start_indices[0] = 0;
    int total_work = 0;
    for (int i = 0; i < t; i++)
    {
        // only consider producers!
        if (i >= cons)
        {
            int workload = (i == cons ? t0_prod_workload : t_prod_workload);
            // n-th thread starts working at index total_work
            start_indices[i] = total_work;
            total_work += workload;
        }
    }
    assert(total_work == items);
    printf("items to add: all integers from 0 to %d stored in %p until %p\n", items - 1, items_to_add, &(items_to_add[items - 1]));
    printf("start indices for producing threads:");
    for (int i = 0; i < t; i++)
        if (i >= cons)
            printf(" %d@t%d", start_indices[i], i);
    printf("\n");
#endif

    // initialize data structure
    init_function(t);

    // pre-populate workload array
    int workloads[t];
    for (int i = 0; i < t; i++)
    {
        int workload = i < cons ? (i == 0 ? t0_cons_workload : t_cons_workload) : (i == cons ? t0_prod_workload : t_prod_workload);
        //printf("thread %2d (is consumer %d) has workload %d\n", i, i < cons, workload);
#if VERIFY
        assert(workload > 0);
#endif
        workloads[i] = workload;
    }

    omp_set_num_threads(t);
    tic = omp_get_wtime();

    {
#pragma omp parallel for shared(global_head_block)
        for (int tid = 0; tid < t; tid++)
        {
            thread_data[tid] = bench_prod_or_cons_thread_function(workloads[tid], tid < cons);
        }
    }

    toc = omp_get_wtime();

    // de-allocate data structure we initialized
    destroy_function();

#if VERIFY
    int sum = 0;
    for (int i = 0; i < items; i++)
        sum += !processed_correctly[i];
    printf("    at the end %d items remained that were not processed correctly\n", sum);
    assert(sum == 0);
    free(start_indices);
    free(items_to_add);
    free(processed_correctly);
#endif

    result.reduced_counters.additions = 0;
    result.reduced_counters.successful_removals = 0;
    result.reduced_counters.failed_removals = 0;
    for (int i = 0; i < t; i++)
    {
        printf("thread %2d has finished with data { .additions=%d, .successful_removals=%d, .failed_removals=%d }\n", i, thread_data[i].additions, thread_data[i].successful_removals, thread_data[i].failed_removals);
        result.reduced_counters.additions += thread_data[i].additions;
        result.reduced_counters.successful_removals += thread_data[i].successful_removals;
        result.reduced_counters.failed_removals += thread_data[i].failed_removals;
    }

    result.time = (toc - tic);
    result.add_time = 0;
    result.try_remove_any_time = 0;
    printf("%s: adding %d elements on %d threads took: %fs\n", name, items, t,
           toc - tic);
    printf(
        "  with %d additions, %d failed removals and %d successful removals: %f removals/s "
        "throughput\n\n",
        result.reduced_counters.additions,
        result.reduced_counters.failed_removals,
        result.reduced_counters.successful_removals,
        result.reduced_counters.successful_removals / result.time);

#if VERIFY
    assert(result.time > 0);
    assert(result.reduced_counters.additions == items);
    assert(result.reduced_counters.successful_removals == items);
#endif

    return result;
}

struct bench_result bench_simple_prod_cons(int t, int items, int cons)
{
    return bench_prod_cons(t, items, cons, "simple", noop_with_int, noop, bench_simple_prod_or_cons_thread);
}

struct bench_result bench_cbag_prod_cons(int t, int items, int cons)
{
    return bench_prod_cons(t, items, cons, "cbag", cbag_init_bag, cbag_destroy_bag, bench_cbag_prod_or_cons_thread);
}

struct bench_result bench_add_try_remove_any(int t, int items, char *name, void (*init_function)(int), void (*destroy_function)(), struct counters (*bench_add_thread_function)(int), struct counters (*bench_try_remove_any_thread_function)(int))
{
    struct bench_result result;
    struct counters thread_data[t];
    double add_tic, add_toc, try_remove_any_tic, try_remove_any_toc;

    if (t < 1 || items < 1)
    {
        printf("invalid bench_add_try_remove_any call: thread count and problem size\n");
        printf("    thread count = %d, problem size = %d\n", t, items);
        return result;
    }

    printf(
        "benchmarking %s add and try_remove_any with %d threads "
        "and problem size %d\n",
        name, t, items);

    // first producer as well as first consumer need to handle excess elements
    int t0_workload = items / t + items % t;
    // all the other take the neatly divisible parts of the work
    int t_workload = items / t;

#if VERIFY
    start_indices = (int *)malloc(sizeof(int) * t);
    items_to_add = (int *)malloc(sizeof(int) * items);
    processed_correctly = (bool *)malloc(sizeof(bool) * items);
    for (int i = 0; i < items; i++)
    {
        items_to_add[i] = i;
        processed_correctly[i] = false;
    }
    // 0th thread starts working (producing) at index 0
    start_indices[0] = 0;
    int total_work = 0;
    for (int i = 0; i < t; i++)
    {
        int workload = (i == 0 ? t0_workload : t_workload);
        // n-th thread starts working at index total_work
        start_indices[i] = total_work;
        total_work += workload;
    }
    assert(total_work == items);
    printf("items to add: all integers from 0 to %d stored in %p until %p\n", items - 1, items_to_add, &(items_to_add[items - 1]));
    printf("start indices for producing threads:");
    for (int i = 0; i < t; i++)
        printf(" %d@t%d", start_indices[i], i);
    printf("\n");
#endif

    // initialize data structure
    init_function(t);

    // pre-populate workload array
    int workloads[t];
    for (int i = 0; i < t; i++)
    {
        int workload = i == 0 ? t0_workload : t_workload;
        // printf("thread %2d has add workload %d\n", i, workload);
#if VERIFY
        assert(workload > 0);
#endif
        workloads[i] = workload;
    }

    omp_set_num_threads(t);
    add_tic = omp_get_wtime();
    {
#pragma omp parallel for shared(global_head_block)
        for (int tid = 0; tid < t; tid++)
        {
            thread_data[tid] = bench_add_thread_function(workloads[tid]);
        }
    }
    add_toc = omp_get_wtime();

    // reduce add results
    result.reduced_counters.additions = 0;
    result.reduced_counters.successful_removals = 0;
    result.reduced_counters.failed_removals = 0;
    for (int i = 0; i < t; i++)
    {
        result.reduced_counters.additions += thread_data[i].additions;
        result.reduced_counters.successful_removals += thread_data[i].successful_removals;
        result.reduced_counters.failed_removals += thread_data[i].failed_removals;
    }

#if VERIFY
    assert(result.reduced_counters.additions == items);
#endif

    try_remove_any_tic = omp_get_wtime();
    {
#pragma omp parallel for shared(global_head_block)
        for (int tid = 0; tid < t; tid++)
        {
            int workload = tid == 0 ? t0_workload : t_workload;
            // printf("thread %2d has try_remove_any workload %d\n", tid, workload);
            thread_data[tid] = bench_try_remove_any_thread_function(workload);
        }
    }
    try_remove_any_toc = omp_get_wtime();

    // de-allocate data structure we initialized
    destroy_function();

#if VERIFY
    int sum = 0;
    for (int i = 0; i < items; i++)
        sum += !processed_correctly[i];
    printf("    at the end %d items remained that were not processed correctly\n", sum);
    assert(sum == 0);
    free(start_indices);
    free(items_to_add);
    free(processed_correctly);
#endif

    // reduce try_remove_any results
    for (int i = 0; i < t; i++)
    {
        printf("thread %2d has finished with data { .additions=%d, .successful_removals=%d, .failed_removals=%d }\n", i, thread_data[i].additions, thread_data[i].successful_removals, thread_data[i].failed_removals);
        result.reduced_counters.additions += thread_data[i].additions;
        result.reduced_counters.successful_removals += thread_data[i].successful_removals;
        result.reduced_counters.failed_removals += thread_data[i].failed_removals;
    }

    result.add_time = (add_toc - add_tic);
    result.try_remove_any_time = (try_remove_any_toc - try_remove_any_tic);
    result.time = result.add_time + result.try_remove_any_time;
    printf("%s: adding %d elements on %d threads took: %fs; removing %d elements took: %fs\n", name, items, t,
           result.add_time, items, result.try_remove_any_time);
    printf(
        "  with %d additions, %d failed removals and %d successful removals: %f removals/s "
        "throughput\n\n",
        result.reduced_counters.additions,
        result.reduced_counters.failed_removals,
        result.reduced_counters.successful_removals,
        result.reduced_counters.successful_removals / result.time);

#if VERIFY
    assert(result.time > 0);
    assert(result.add_time > 0);
    assert(result.try_remove_any_time > 0);
    assert(result.reduced_counters.successful_removals == items);
#endif

    return result;
}

struct bench_result bench_simple_add_try_remove_any(int t, int items)
{
    return bench_add_try_remove_any(t, items, "simple", noop_with_int, noop, bench_simple_add_thread, bench_simple_try_remove_any_thread);
}

struct bench_result bench_cbag_add_try_remove_any(int t, int items)
{
    return bench_add_try_remove_any(t, items, "cbag", cbag_init_bag, cbag_destroy_bag, bench_cbag_add_thread, bench_cbag_try_remove_any_thread);
}

/* main is not relevant for benchmark.py but necessary when run alone for
 * testing.
 */
int main(int argc, char *argv[])
{
    (void)argc;
    (void)argv;
    // printf("max_lluint %llx mask_pointer %llx mask_mark1 %llx mask_mark2 %llx\n", MAX_LLUINT, MASK_POINTER, MASK_MARK1, MASK_MARK2);

    // markrefs utility function tests
    void *maxp = MAX_POINTER;
    printf("maxp %p, ref %p, mark1 %i, mark2 %i\n", maxp, block_get_ref(maxp), block_get_mark1(maxp), block_get_mark2(maxp));
    printf("maxp %p, ref %p, mark1 %i, mark2 %i\n", NULL, block_get_ref(NULL), block_get_mark1(NULL), block_get_mark2(NULL));
    maxp = block_get_ref(maxp);
    printf("maxp %p, ref %p, mark1 %i, mark2 %i\n", maxp, block_get_ref(maxp), block_get_mark1(maxp), block_get_mark2(maxp));
    block_set_mark1((block_t **)&maxp, true);
    block_set_mark2((block_t **)&maxp, true);
    printf("maxp %p, ref %p, mark1 %i, mark2 %i\n", maxp, block_get_ref(maxp), block_get_mark1(maxp), block_get_mark2(maxp));

    // this is the benchmark
    int workload = 10;
    for (int k = 0; k < 4; k++)
    {
        workload *= 10;
        for (int t = 2; t < 32; t *= 2)
        {
            printf("starting test t=%d and n=%d\n", t, workload);
            bench_simple_prod_cons(t, workload, t / 2);
            bench_cbag_prod_cons(t, workload, t / 2);
            bench_simple_prod_cons(t, workload, 1);
            bench_cbag_prod_cons(t, workload, 1);
            bench_simple_prod_cons(t, workload, t - 1);
            bench_cbag_prod_cons(t, workload, t - 1);
            bench_simple_add_try_remove_any(t, workload);
            bench_cbag_add_try_remove_any(t, workload);
        }
    }
}
