#include <stdio.h>
#include <stdlib.h>
#include <omp.h>

#define BUCKETS 10
#define MAX_N 2000000

/* A bucket holds the values from one range, in no particular order. */
typedef struct {
    int count;
    int items[MAX_N];
    omp_lock_t lock;
} bucket_t;

bucket_t buckets[BUCKETS];

/* A made-up value from 0 to 999 for item i, with some work to compute it. */
int value(int i) {
    unsigned int x = (unsigned int)i;
    for (int step = 0; step < 50; step++) {
        x = x * 1103515245u + 12345u;
    }
    return (int)((x >> 8) % 1000);
}

/* Two steps: store the value, then count it. */
void append(bucket_t *b, int v) {
    b->items[b->count] = v;
    b->count++;
}

void empty_buckets(void) {
    for (int b = 0; b < BUCKETS; b++) {
        buckets[b].count = 0;
    }
}

/* Check each bucket's count and sum against a serial calculation. */
int check(int n) {
    int expected_count[BUCKETS] = {0};
    long long expected_sum[BUCKETS] = {0};
    for (int i = 0; i < n; i++) {
        int v = value(i);
        expected_count[v / 100]++;
        expected_sum[v / 100] += v;
    }

    int total = 0;
    for (int b = 0; b < BUCKETS; b++) {
        long long sum = 0;
        for (int k = 0; k < buckets[b].count; k++) {
            sum += buckets[b].items[k];
        }
        if (buckets[b].count != expected_count[b] || sum != expected_sum[b]) {
            return 0;
        }
        total += buckets[b].count;
    }
    return total == n;
}

int main(int argc, char *argv[]) {
    int thread_count = 4;
    int n = MAX_N;
    if (argc > 1) {
        thread_count = (int)strtol(argv[1], NULL, 10);
    }
    if (argc > 2) {
        n = (int)strtol(argv[2], NULL, 10);
    }
    if (n < 1 || n > MAX_N) {
        fprintf(stderr, "n must be from 1 through %d\n", MAX_N);
        return 1;
    }
    for (int b = 0; b < BUCKETS; b++) {
        omp_init_lock(&buckets[b].lock);
    }
    printf("Threads: %d   Values: %d\n", thread_count, n);

    /* Version 1: one critical section protects every bucket. */
    empty_buckets();
    double start = omp_get_wtime();
#   pragma omp parallel for num_threads(thread_count)
    for (int i = 0; i < n; i++) {
        int v = value(i);
#       pragma omp critical
        append(&buckets[v / 100], v);
    }
    double elapsed = omp_get_wtime() - start;
    printf("One critical section: %.3f seconds, %s\n", elapsed,
           check(n) ? "correct" : "WRONG");

    /* Version 2: each bucket has its own lock. */
    empty_buckets();
    start = omp_get_wtime();
#   pragma omp parallel for num_threads(thread_count)
    for (int i = 0; i < n; i++) {
        int v = value(i);
        bucket_t *b = &buckets[v / 100];
        omp_set_lock(&b->lock);
        append(b, v);
        omp_unset_lock(&b->lock);
    }
    elapsed = omp_get_wtime() - start;
    printf("A lock per bucket:    %.3f seconds, %s\n", elapsed,
           check(n) ? "correct" : "WRONG");

    for (int b = 0; b < BUCKETS; b++) {
        omp_destroy_lock(&buckets[b].lock);
    }
    return 0;
}