#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <time.h>
#include <pthread.h>
#include <math.h>
#include <string.h>

static pthread_barrier_t sync_barrier;
bool solved;
bool allow_print;

// Struct containing all the useful information for a thread
typedef struct thread_data {
    unsigned long start_row;
    unsigned long end_row;
    double **array;
    unsigned long dimension;
    double precision;
    bool finished;
} THREAD_DATA_T;

void printSquare(double **array, unsigned long dimension) {
    unsigned long i, j;
    for (i = 0; i < dimension; i++) {
        for (j = 0; j < dimension; j++) printf("%f ", array[i][j]);
        printf("\n");
    }
}

long double time_seconds(struct timespec begin, struct timespec end) {
    long long sec_diff = end.tv_sec - begin.tv_sec;
    long long nsec_diff = end.tv_nsec - begin.tv_nsec;
    long long as_ns = sec_diff * 1000000000L + nsec_diff;
    return (long double) as_ns / 1000000000.0;
}

void solveArray(double **array, unsigned long dimension, double precision,
                unsigned long start, unsigned long end) {
    unsigned long i, j;
    for (i = start; i < end; i++) {
        for (j = 1; j < dimension - 1; j++) {
            // no need to gain locks on any of the array elements as we are the
            // only one writing to them
            // and whether we get the old or new values of an element
            // doesn't matter
            double previous = array[i][j];
            double avg = (array[i][j - 1] + array[i][j + 1] + array[i - 1][j] +
                          array[i + 1][j]) / 4.0;
            array[i][j] = avg;
            // only do the actual check if we haven't already failed
            if (solved == true && fabs(avg - previous) > precision) {
                // dont bother getting a lock here because the only other value
                // it will be set to
                // is false and we don't care if it gets messed up
                solved = false;
            }
        }
    }
}

void *bootstrap_thread(void *arg) {
    THREAD_DATA_T *td = (THREAD_DATA_T *) arg;
    // Each thread will do a single iteration over the array before waiting for
    // every thread to finish
    // and then waiting for the main thread to check if things are solved
    while (1) {
        if (td->finished) {
            // this is only here to stop static code analysis
            // whining about infinite loops
            return NULL;
        }
        solveArray(td->array, td->dimension, td->precision, td->start_row,
                   td->end_row);
        // first barrier waits for all threads to finish
        pthread_barrier_wait(&sync_barrier);
        // wait for main thread to check completion
        pthread_barrier_wait(&sync_barrier);
    }
    return NULL;
}

double **constructBigArray(double top, double bottom, double left, double right,
                           long unsigned int n) {
    // here lies exciting memory allocation and array initialisation, boring..
    long unsigned int i, j, k;
    double **arr = (double **) malloc(n * sizeof(double *));
    for (k = 0; k < n; k++) {
        arr[k] = (double *) malloc(n * sizeof(double));
    }
    for (i = 0; i < n; i++) {
        for (j = 0; j < n; j++) {
            if (i == 0) {
                if (j == 0 || j == n - 1) arr[i][j] = 0;
                else arr[i][j] = top;
            } else if (i == n - 1) {
                if (j == 0 || j == n - 1) arr[i][j] = 0;
                else arr[i][j] = bottom;
            } else if (j == 0) {
                arr[i][j] = left;
            } else if (j == n - 1) {
                arr[i][j] = right;
            } else {
                arr[i][j] = 0;
            }
        }
    }
    return arr;
}

// CHANGE ME to set thread count
#define NUM_THREADS 10

// struct to hold some argument data
typedef struct arg_data {
    double precision;
    unsigned long dimension;
    double top_start;
    double bottom_start;
    double left_start;
    double right_start;
} ARG_DATA_T;

// helper function for reading in arguments
int argsContain(const char *val, char **arr, int size) {
    int i;
    for (i = 0; i < size; i++) {
        if (strcmp(arr[i], val) == 0)
            return i;
    }
    return -1;
}

// take a look at argv and print out usage or take in some arguments
ARG_DATA_T *parseArgs(int argc, char **argv) {
    int tmp;
    ARG_DATA_T *args = (ARG_DATA_T *) malloc(sizeof(ARG_DATA_T));
    allow_print = true;
    if ((tmp = argsContain("-p", argv, argc)) >= 0 ||
        (tmp = argsContain("--precision", argv, argc)) >= 0) {
        args->precision = strtod(argv[tmp + 1], NULL);
    } else {
        args->precision = 0.00001;
    }
    if ((tmp = argsContain("-d", argv, argc)) >= 0 ||
        (tmp = argsContain("--dimension", argv, argc)) >= 0) {
        args->dimension = (unsigned long) strtol(argv[tmp + 1], NULL, 10);
    } else {
        args->dimension = 10;
    }
    if ((tmp = argsContain("-t", argv, argc)) >= 0 ||
        (tmp = argsContain("--top", argv, argc)) >= 0) {
        args->top_start = strtod(argv[tmp + 1], NULL);
    } else {
        args->top_start = 1;
    }
    if ((tmp = argsContain("-l", argv, argc)) >= 0 ||
        (tmp = argsContain("--left", argv, argc)) >= 0) {
        args->left_start = strtod(argv[tmp + 1], NULL);
    } else {
        args->left_start = 2;
    }
    if ((tmp = argsContain("-r", argv, argc)) >= 0 ||
        (tmp = argsContain("--right", argv, argc)) >= 0) {
        args->right_start = strtod(argv[tmp + 1], NULL);
    } else {
        args->right_start = 4;
    }
    if ((tmp = argsContain("-b", argv, argc)) >= 0 ||
        (tmp = argsContain("--bottom", argv, argc)) >= 0) {
        args->bottom_start = strtod(argv[tmp + 1], NULL);
    } else {
        args->bottom_start = 3;
    }
    if ((tmp = argsContain("-q", argv, argc)) > 0 ||
        (tmp = argsContain("--quiet", argv, argc)) >= 0) {
      allow_print = false;
    }
    if (argsContain("-h", argv, argc) >= 0 ||
        argsContain("--help", argv, argc) >= 0) {
        printf("Usage: %s [-p N] [-d N] [-t N] [-l N] [-r N] [-b N]\n",
               argv[0]);
        printf("Constructs a 2D array of doubles, with values from arguments "
                       "of size dimension^2 and performs relaxation method"
                       " to given precision.\n");
        printf("Options:\n");
        printf("-p/--precision: A double value to specify precision. "
                       "Default: 0.0001\n");
        printf("-d/--dimension: A long value to specify width and height of "
                       "grid. Default: 10\n");
        printf("-t/--top: A double indicating the values along the top of the "
                       "array. Default: 1.0\n");
        printf("-l/--left: A double indicating the values along the left side"
                       " of the array. Default: 2.0\n");
        printf("-b/--bottom: A double indicating the values along the bottom "
                       "of the array. Default: 3.0\n");
        printf("-r/--right: A double indicating the values along the right "
                       "side of the array. Default: 4.0\n");
        printf("-q/--quiet: Disables all printing except the result.\n");
        exit(0);
    }
    return args;
}

// this function is used to return the number of rows a given thread will work
// on
// for an amount of rows
unsigned long
spread(unsigned long group, unsigned long totalGroups, unsigned long n) {
    if (totalGroups % n == 0) return totalGroups / n;
    else {
        unsigned long tmp = (totalGroups / n);
        if (totalGroups % n >= group) return tmp + 1;
        else return tmp;
    }
}

// helper function as not everyone defines the MIN macro
unsigned long mini(unsigned long a, unsigned long b) {
    if (a < b) {
        return a;
    } else {
        return b;
    }
}

int main(int argc, char **argv) {
    int rc;
    pthread_t thread[NUM_THREADS];
    THREAD_DATA_T *td[NUM_THREADS];
    ARG_DATA_T *args;
    struct timespec begin, end;

    args = parseArgs(argc, argv);

    // the barrier should be at least the array size, subtract 2, as we won't
    // use all threads if there are not enough single rows to share out
    pthread_barrier_init(&sync_barrier, NULL,
                         (unsigned int) mini(args->dimension - 2, NUM_THREADS));
    double **array = constructBigArray(args->top_start, args->bottom_start,
                                       args->left_start, args->right_start,
                                       args->dimension);

    // use clock_gettime to measure performance as it has best resolution
    clock_gettime(CLOCK_MONOTONIC, &begin);
    td[0] = (THREAD_DATA_T *) malloc(sizeof(THREAD_DATA_T));
    solved = true;

    // create NUM_THREADS-1 pthreads, starting from 1 - because the main thread
    // is one of our threads
    unsigned long rowCounter = 1 + spread(1, args->dimension - 2, NUM_THREADS);
    for (unsigned long i = 1; i < NUM_THREADS; i++) {
        // divides up the array into rows, so that each thread is operating on
        // chunks of contiguous memory
        unsigned long start_row, end_row;
        start_row = rowCounter;
        rowCounter += spread(i + 1, args->dimension - 2, NUM_THREADS);
        end_row = rowCounter;

        td[i] = (THREAD_DATA_T *) malloc(sizeof(THREAD_DATA_T));
        td[i]->start_row = start_row;
        td[i]->end_row = end_row;
        td[i]->array = array;
        td[i]->dimension = args->dimension;
        td[i]->precision = args->precision;

        td[i]->finished = false;

        if (end_row - start_row > 0) { // only create threads that will do work
            if ((rc = pthread_create(&thread[i], NULL, bootstrap_thread,
                                     (void *) td[i]))) {
                printf("failed to create threads: %d\n", rc);
                return 1;
            }
        }
    }
    while (1) {
        //main thread should also be solving
        solveArray(array, args->dimension, args->precision, 1,
                   1 + spread(1, args->dimension - 2, NUM_THREADS));
        pthread_barrier_wait(&sync_barrier);
        if (solved) {
            clock_gettime(CLOCK_MONOTONIC, &end);
            if (allow_print) printf("took me %Lfs\n", time_seconds(begin, end));
            for (int i = 1; i < NUM_THREADS; i++) {
                td[i]->finished = true;
            }
            if (allow_print) printf("====== SOLVED =====\n");
            printSquare(array, args->dimension);
            return 0;
        }
        solved = true; // make sure to reset this
        pthread_barrier_wait(&sync_barrier);
    }

}