#include <stdio.h>
#include <mpi.h>
#include <assert.h>
#include <string.h>
#include "../dispatch.h"

#define INDEX_2D_TO_1D(y, x, nb_columns) ((y) * (nb_columns) + (x))
#define INDEX_3D_TO_1D(y, x, z, nb_columns, nb_depths) ((x) + (nb_columns) * ((y) + (nb_depths) * (z)))

void init_chunk_data(chunk_info_t *ci) {
    for (int y = 0; y < ci->dimensions[0]; ++y) {
        for (int x = 0; x < ci->dimensions[1]; ++x) {
            for (int z = 0; z < ci->dimensions[2]; ++z) {
                ((int *) ci->data)[INDEX_3D_TO_1D(y, x, z, ci->dimensions[1], ci->dimensions[2])] = y + x + z;
            }
        }
    }
}

void tests_1_process(struct futhark_context *fc) {
    printf("Testing dispatch_context_new with 1 process. Dimensions = 1x15...\n");
    {
        int dimensions[1] = {15};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 1);
        chunk_info_t ci = get_chunk_info(dc);
        assert(ci.dimensions[0] == 1 && ci.dimensions[1] == 15);
        assert(ci.data != NULL);
        assert(ci.count == 15);
        assert(ci.y == 0 && ci.x == 0);

        init_chunk_data(&ci);

        printf("\tTesting get_inner_envelope. Thickness = 1...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 1);
            int *data = (int *) ci.data;
            int *west = inner_envelope->front->west->data;
            int *east = inner_envelope->front->east->data;
            assert(west[0] == data[0]);
            assert(east[0] == data[14]);
            envelope_free(inner_envelope);
        }

        printf("\tTesting get_inner_envelope. Thickness = 2...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 2);
            int *west = inner_envelope->front->west->data;
            int *east = inner_envelope->front->east->data;
            assert(west[0] == 0 && west[1] == 1);
            assert(east[0] == 13 && east[1] == 14);
            envelope_free(inner_envelope);

            printf("\tTesting get_inner_envelope. Thickness = 16...\n");
            inner_envelope = get_inner_envelope(dc, fc, 16);
            assert(memcmp(inner_envelope->front->west->data, ci.data, (size_t) ci.count * sizeof(int)) == 0);
            assert(memcmp(inner_envelope->front->east->data, ci.data, (size_t) ci.count * sizeof(int)) == 0);
            envelope_free(inner_envelope);
        }

        printf("\tTesting get_outer_envelope. Thickness = 1...\n");
        {
            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 1);
            int *data = (int *) ci.data;
            int *west = outer_envelope->front->west->data;
            int *east = outer_envelope->front->east->data;
            assert(west[0] == data[14]);
            assert(east[0] == data[0]);
            envelope_free(outer_envelope);
        }

        printf("\tTesting get_outer_envelope. Thickness = 2...\n");
        {
            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 2);
            int *data = (int *) ci.data;
            int *west = outer_envelope->front->west->data;
            int *east = outer_envelope->front->east->data;
            assert(west[0] == data[13] && west[1] == data[14]);
            assert(east[0] == data[0] && east[1] == data[1]);
            envelope_free(outer_envelope);
        }

        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 1 process. Dimensions = 10x15...\n");
    {
        int dimensions[2] = {10, 15};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 15);
        assert(ci.data != NULL);
        assert(ci.count == 150);
        assert(ci.y == 0 && ci.x == 0);

        init_chunk_data(&ci);

        printf("\tTesting get_inner_envelope. Thickness = 1...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 1);

            int expected_north[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14};
            int expected_east[] = {14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_south[] = {9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_west[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
            int expected_nw[] = {0};
            int expected_ne[] = {14};
            int expected_se[] = {23};
            int expected_sw[] = {9};

            assert(memcmp(expected_north, inner_envelope->front->north->data, sizeof(expected_north)) == 0);
            assert(memcmp(expected_east, inner_envelope->front->east->data, sizeof(expected_east)) == 0);
            assert(memcmp(expected_south, inner_envelope->front->south->data, sizeof(expected_south)) == 0);
            assert(memcmp(expected_west, inner_envelope->front->west->data, sizeof(expected_west)) == 0);
            assert(memcmp(expected_nw, inner_envelope->front->north_west->data, sizeof(expected_nw)) == 0);
            assert(memcmp(expected_ne, inner_envelope->front->north_east->data, sizeof(expected_ne)) == 0);
            assert(memcmp(expected_se, inner_envelope->front->south_east->data, sizeof(expected_se)) == 0);
            assert(memcmp(expected_sw, inner_envelope->front->south_west->data, sizeof(expected_sw)) == 0);

            envelope_free(inner_envelope);
        }

        printf("\tTesting get_outer_envelope. Thickness = 1...\n");
        {
            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 1);

            int expected_south[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14};
            int expected_west[] = {14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_north[] = {9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_east[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
            int expected_se[] = {0};
            int expected_sw[] = {14};
            int expected_nw[] = {23};
            int expected_ne[] = {9};

            assert(memcmp(expected_north, outer_envelope->front->north->data, sizeof(expected_north)) == 0);
            assert(memcmp(expected_east, outer_envelope->front->east->data, sizeof(expected_east)) == 0);
            assert(memcmp(expected_south, outer_envelope->front->south->data, sizeof(expected_south)) == 0);
            assert(memcmp(expected_west, outer_envelope->front->west->data, sizeof(expected_west)) == 0);
            assert(memcmp(expected_nw, outer_envelope->front->north_west->data, sizeof(expected_nw)) == 0);
            assert(memcmp(expected_ne, outer_envelope->front->north_east->data, sizeof(expected_ne)) == 0);
            assert(memcmp(expected_se, outer_envelope->front->south_east->data, sizeof(expected_se)) == 0);
            assert(memcmp(expected_sw, outer_envelope->front->south_west->data, sizeof(expected_sw)) == 0);

            envelope_free(outer_envelope);
        }

        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 1 process. Dimensions = 10x15x1...\n");
    {
        int dimensions[3] = {10, 15, 1};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 3);
        chunk_info_t ci = get_chunk_info(dc);
        assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 15 && ci.dimensions[2] == 1);
        assert(ci.data != NULL);
        assert(ci.count == 150);
        assert(ci.y == 0 && ci.x == 0 && ci.z == 0);

        init_chunk_data(&ci);

        printf("\tTesting get_inner_envelope. Thickness = 1...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 1);
            int expected_north[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14};
            int expected_east[] = {14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_south[] = {9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_west[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
            int expected_nw[] = {0};
            int expected_ne[] = {14};
            int expected_se[] = {23};
            int expected_sw[] = {9};

            assert(memcmp(expected_east, inner_envelope->front->east->data, sizeof(expected_east)) == 0);
            assert(memcmp(expected_west, inner_envelope->front->west->data, sizeof(expected_west)) == 0);

            envelope_free(inner_envelope);
        }

        printf("\tTesting get_outer_envelope. Thickness = 1...\n");
        {
            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 1);

            int expected_south[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14};
            int expected_west[] = {14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_north[] = {9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23};
            int expected_east[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
            int expected_se[] = {0};
            int expected_sw[] = {14};
            int expected_nw[] = {23};
            int expected_ne[] = {9};

            assert(memcmp(expected_north, outer_envelope->top->north->data, sizeof(expected_north)) == 0);
            assert(memcmp(expected_east, outer_envelope->front->east->data, sizeof(expected_east)) == 0);
            assert(memcmp(expected_south, outer_envelope->bottom->south->data, sizeof(expected_south)) == 0);
            assert(memcmp(expected_west, outer_envelope->front->west->data, sizeof(expected_west)) == 0);
            assert(memcmp(expected_nw, outer_envelope->top->north_west->data, sizeof(expected_nw)) == 0);
            assert(memcmp(expected_ne, outer_envelope->top->north_east->data, sizeof(expected_ne)) == 0);
            assert(memcmp(expected_se, outer_envelope->bottom->south_east->data, sizeof(expected_se)) == 0);
            assert(memcmp(expected_sw, outer_envelope->bottom->south_west->data, sizeof(expected_sw)) == 0);

            envelope_free(outer_envelope);
        }

        dispatch_context_free(dc);
    }
}

void tests_2_processes(int my_rank, struct futhark_context *fc) {
    printf("Testing dispatch_context_new with 2 processes. Dimensions = 1x16...\n");
    {
        int dimensions[1] = {16};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 1);
        chunk_info_t ci = get_chunk_info(dc);

        if (my_rank == 0) {
            assert(ci.dimensions[0] == 1 && ci.dimensions[1] == 8);
            assert(ci.data != NULL);
            assert(ci.count == 8);
            assert(ci.y == 0 && ci.x == 0);
        } else {
            assert(ci.dimensions[0] == 1 && ci.dimensions[1] == 8);
            assert(ci.data != NULL);
            assert(ci.count == 8);
            assert(ci.y == 0 && ci.x == 8);
        }

        init_chunk_data(&ci);
        printf("\tTesting get_inner_envelope. Thickness = 1...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 1);
            int *data = (int *) ci.data;
            int *west = inner_envelope->front->west->data;
            int *east = inner_envelope->front->east->data;
            assert(west[0] == data[0]);
            assert(east[0] == data[ci.dimensions[1] - 1]);
            envelope_free(inner_envelope);
        }

        printf("\tTesting get_outer_envelope. Thickness = 9...\n");
        {
            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 9);
            int expected_west[] = {0, 1, 2, 3, 4, 5, 6, 7};
            int expected_east[] = {0, 1, 2, 3, 4, 5, 6, 7};

            assert(memcmp(expected_east, outer_envelope->front->east->data, sizeof(expected_east)) == 0);
            assert(memcmp(expected_west, outer_envelope->front->west->data, sizeof(expected_west)) == 0);

            envelope_free(outer_envelope);
        }

        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 2 processes. Dimensions = 1x15...\n");
    {
        int dimensions[1] = {15};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 1);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 1 && ci.dimensions[1] == 8);
            assert(ci.data != NULL);
            assert(ci.count == 8);
            assert(ci.y == 0 && ci.x == 0);
        } else {
            assert(ci.dimensions[0] == 1 && ci.dimensions[1] == 7);
            assert(ci.data != NULL);
            assert(ci.count == 7);
            assert(ci.y == 0 && ci.x == 8);
        }

        init_chunk_data(&ci);

        printf("\tTesting get_inner_envelope. Thickness = 1...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 1);
            int *data = (int *) ci.data;
            int *west = inner_envelope->front->west->data;
            int *east = inner_envelope->front->east->data;
            assert(west[0] == data[0]);
            assert(east[0] == data[ci.dimensions[1] - 1]);
            envelope_free(inner_envelope);
        }

        printf("\tTesting get_inner_envelope. Thickness = 8...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 8);

            assert(memcmp(inner_envelope->front->west->data, ci.data, (size_t) ci.count * sizeof(int)) == 0);
            assert(memcmp(inner_envelope->front->east->data, ci.data, (size_t) ci.count * sizeof(int)) == 0);

            envelope_free(inner_envelope);
        }

        printf("\tTesting get_outer_envelope. Thickness = 9...\n");
        {
            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 9);
            if (my_rank == 0) {
                int expected_west[] = {0, 1, 2, 3, 4, 5, 6};
                int expected_east[] = {0, 1, 2, 3, 4, 5, 6};

                assert(memcmp(expected_east, outer_envelope->front->east->data, sizeof(expected_east)) == 0);
                assert(memcmp(expected_west, outer_envelope->front->west->data, sizeof(expected_west)) == 0);
            } else {
                int expected_west[] = {0, 1, 2, 3, 4, 5, 6, 7};
                int expected_east[] = {0, 1, 2, 3, 4, 5, 6, 7};
            }
            envelope_free(outer_envelope);
        }

        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 2 processes. Dimensions = 20x40...\n");
    {
        int dimensions[2] = {20, 40};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 20 && ci.dimensions[1] == 20);
            assert(ci.data != NULL);
            assert(ci.count == 400);
            assert(ci.y == 0 && ci.x == 0);
        } else {
            assert(ci.dimensions[0] == 20 && ci.dimensions[1] == 20);
            assert(ci.data != NULL);
            assert(ci.count == 400);
            assert(ci.y == 0 && ci.x == 20);
        }

        init_chunk_data(&ci);

        printf("\tTesting get_inner_envelope. Thickness = 4...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 4);
            if (my_rank == 0) {
                int expected_north[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 1, 2, 3,
                                        4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 2, 3, 4, 5, 6, 7,
                                        8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 3, 4, 5, 6, 7, 8, 9, 10,
                                        11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22
                };

                int expected_south[] = {16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
                                        17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36,
                                        18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37,
                                        19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38
                };

                int expected_west[] = {0, 1, 2, 3, 1, 2, 3, 4, 2, 3, 4, 5, 3, 4, 5, 6, 4, 5, 6, 7, 5, 6, 7, 8, 6, 7, 8,
                                       9, 7, 8, 9, 10, 8, 9, 10, 11, 9, 10, 11, 12, 10, 11, 12, 13, 11, 12, 13, 14, 12,
                                       13, 14, 15, 13, 14, 15, 16, 14, 15, 16, 17, 15, 16, 17, 18, 16, 17, 18, 19, 17,
                                       18, 19, 20, 18, 19, 20, 21,
                                       19, 20, 21, 22};

                int expected_east[] = {
                        16, 17, 18, 19, 17, 18, 19, 20, 18, 19, 20, 21, 19, 20, 21, 22, 20, 21, 22, 23,
                        21, 22, 23, 24, 22, 23, 24, 25, 23, 24, 25, 26, 24, 25, 26, 27, 25, 26, 27, 28,
                        26, 27, 28, 29, 27, 28, 29, 30, 28, 29, 30, 31, 29, 30, 31, 32, 30, 31, 32, 33,
                        31, 32, 33, 34, 32, 33, 34, 35, 33, 34, 35, 36, 34, 35, 36, 37, 35, 36, 37, 38,
                };

                int expected_nw[] = {0, 1, 2, 3,
                                     1, 2, 3, 4,
                                     2, 3, 4, 5,
                                     3, 4, 5, 6};

                int expected_sw[] = {16, 17, 18, 19,
                                     17, 18, 19, 20,
                                     18, 19, 20, 21,
                                     19, 20, 21, 22};

                int expected_se[] = {32, 33, 34, 35,
                                     33, 34, 35, 36,
                                     34, 35, 36, 37,
                                     35, 36, 37, 38};

                int expected_ne[] = {
                        16, 17, 18, 19,
                        17, 18, 19, 20,
                        18, 19, 20, 21,
                        19, 20, 21, 22
                };

                assert(memcmp(expected_north, inner_envelope->front->north->data, sizeof(expected_north)) == 0);
                assert(memcmp(expected_east, inner_envelope->front->east->data, sizeof(expected_east)) == 0);
                assert(memcmp(expected_south, inner_envelope->front->south->data, sizeof(expected_south)) == 0);
                assert(memcmp(expected_west, inner_envelope->front->west->data, sizeof(expected_west)) == 0);
                assert(memcmp(expected_nw, inner_envelope->front->north_west->data, sizeof(expected_nw)) == 0);
                assert(memcmp(expected_ne, inner_envelope->front->north_east->data, sizeof(expected_ne)) == 0);
                assert(memcmp(expected_se, inner_envelope->front->south_east->data, sizeof(expected_se)) == 0);
                assert(memcmp(expected_sw, inner_envelope->front->south_west->data, sizeof(expected_sw)) == 0);

            }
            envelope_free(inner_envelope);
        }

        printf("\tTesting get_outer_envelope. Thickness = 4...\n");
        {
            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 4);


            envelope_free(outer_envelope);
        }

        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 2 processes. Dimensions = 21x41...\n");
    {
        int dimensions[2] = {21, 41};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 21 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 441);
            assert(ci.y == 0 && ci.x == 0);
        } else {
            assert(ci.dimensions[0] == 21 && ci.dimensions[1] == 20);
            assert(ci.data != NULL);
            assert(ci.count == 420);
            assert(ci.y == 0 && ci.x == 21);
        }
        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 2 processes. Dimensions = 21x41x3...\n");
    {
        int dimensions[3] = {21, 41, 3};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 3);
        dispatch_context_print(dc);
        chunk_info_t ci = get_chunk_info(dc);
        chunk_info_print(&ci);

        envelope_t *outer_envelope = get_outer_envelope(dc, fc, 1);
        envelope_free(outer_envelope);

        dispatch_context_free(dc);
    }
}

void tests_3_processes(int my_rank, struct futhark_context *fc) {
    printf("Testing dispatch_context_new with 3 processes. Dimensions = 20x40...\n");
    {
        int dimensions[2] = {20, 40};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 20 && ci.dimensions[1] == 14);
            assert(ci.data != NULL);
            assert(ci.count == 280);
            assert(ci.y == 0 && ci.x == 0);
        } else if (my_rank == 1) {
            assert(ci.dimensions[0] == 20 && ci.dimensions[1] == 13);
            assert(ci.data != NULL);
            assert(ci.count == 260);
            assert(ci.y == 0 && ci.x == 14);
        } else {
            assert(ci.dimensions[0] == 20 && ci.dimensions[1] == 13);
            assert(ci.data != NULL);
            assert(ci.count == 260);
            assert(ci.y == 0 && ci.x == 27);
        }
        dispatch_context_free(dc);
    }
    printf("Testing dispatch_context_new with 2 processes. Dimensions = 21x42...\n");
    {
        int dimensions[2] = {21, 42};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 21 && ci.dimensions[1] == 14);
            assert(ci.data != NULL);
            assert(ci.count == 294);
            assert(ci.y == 0 && ci.x == 0);
        } else if (my_rank == 1) {
            assert(ci.dimensions[0] == 21 && ci.dimensions[1] == 14);
            assert(ci.data != NULL);
            assert(ci.count == 294);
            assert(ci.y == 0 && ci.x == 14);
        } else {
            assert(ci.dimensions[0] == 21 && ci.dimensions[1] == 14);
            assert(ci.data != NULL);
            assert(ci.count == 294);
            assert(ci.y == 0 && ci.x == 28);
        }
        dispatch_context_free(dc);
    }
}

void tests_4_processes(int my_rank, struct futhark_context *fc) {
    printf("Testing dispatch_context_new with 4 processes. Dimensions = 20x40...\n");
    {
        int dimensions[2] = {20, 40};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 20);
            assert(ci.data != NULL);
            assert(ci.count == 200);
            assert(ci.y == 0 && ci.x == 0);
        } else if (my_rank == 1) {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 20);
            assert(ci.data != NULL);
            assert(ci.count == 200);
            assert(ci.y == 0 && ci.x == 20);
        } else if (my_rank == 2) {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 20);
            assert(ci.data != NULL);
            assert(ci.count == 200);
            assert(ci.y == 10 && ci.x == 0);
        } else {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 20);
            assert(ci.data != NULL);
            assert(ci.count == 200);
            assert(ci.y == 10 && ci.x == 20);
        }

        printf("\tTesting get_inner_envelope. Thickness = 4...\n");
        {
            envelope_t *inner_envelope = get_inner_envelope(dc, fc, 4);
            envelope_free(inner_envelope);

            envelope_t *outer_envelope = get_outer_envelope(dc, fc, 4);
            envelope_free(outer_envelope);
        }

        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 4 processes. Dimensions = 21x42...\n");
    {
        int dimensions[2] = {21, 42};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 11 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 231);
            assert(ci.y == 0 && ci.x == 0);
        } else if (my_rank == 1) {
            assert(ci.dimensions[0] == 11 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 231);
            assert(ci.y == 0 && ci.x == 21);
        } else if (my_rank == 2) {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 210);
            assert(ci.y == 11 && ci.x == 0);
        } else {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 210);
            assert(ci.y == 11 && ci.x == 21);
        }
        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 4 processes. Dimensions = 21x42...\n");
    {
        int dimensions[2] = {21, 42};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);
        if (my_rank == 0) {
            assert(ci.dimensions[0] == 11 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 231);
            assert(ci.y == 0 && ci.x == 0);
        } else if (my_rank == 1) {
            assert(ci.dimensions[0] == 11 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 231);
            assert(ci.y == 0 && ci.x == 21);
        } else if (my_rank == 2) {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 210);
            assert(ci.y == 11 && ci.x == 0);
        } else {
            assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 21);
            assert(ci.data != NULL);
            assert(ci.count == 210);
            assert(ci.y == 11 && ci.x == 21);
        }
        dispatch_context_free(dc);
    }
}

void tests_6_processes(int my_rank, struct futhark_context *fc) {
    printf("Testing dispatch_context_new with 6 processes. Dimensions = 20x40...\n");
    {
        int dimensions[2] = {20, 40};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 2);
        chunk_info_t ci = get_chunk_info(dc);

        switch (my_rank) {
            case 0:
                assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 14);
                assert(ci.data != NULL);
                assert(ci.count == 140);
                assert(ci.y == 0 && ci.x == 0);
                break;
            case 1:
                assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 13);
                assert(ci.data != NULL);
                assert(ci.count == 130);
                assert(ci.y == 0 && ci.x == 14);
                break;
            case 2:
                assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 13);
                assert(ci.data != NULL);
                assert(ci.count == 130);
                assert(ci.y == 0 && ci.x == 27);
                break;
            case 3:
                assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 14);
                assert(ci.data != NULL);
                assert(ci.count == 140);
                assert(ci.y == 10 && ci.x == 0);
                break;
            case 4:
                assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 13);
                assert(ci.data != NULL);
                assert(ci.count == 130);
                assert(ci.y == 10 && ci.x == 14);
                break;
            case 5:
                assert(ci.dimensions[0] == 10 && ci.dimensions[1] == 13);
                assert(ci.data != NULL);
                assert(ci.count == 130);
                assert(ci.y == 10 && ci.x == 27);
                break;
            default:
                break;
        }
        envelope_t *inner_envelope = get_inner_envelope(dc, fc, 4);
        envelope_free(inner_envelope);
        envelope_t *outer_envelope = get_outer_envelope(dc, fc, 4);
        envelope_free(outer_envelope);
        dispatch_context_free(dc);
    }

    printf("Testing dispatch_context_new with 6 processes. Dimensions = 21x42x5...\n");
    {
        int dimensions[3] = {24, 24, 6};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_UINT8_T, 3);
        dispatch_context_print(dc);
        chunk_info_t ci = get_chunk_info(dc);

        envelope_t *inner_envelope = get_inner_envelope(dc, fc, 1);
        envelope_free(inner_envelope);

        envelope_t *outer_envelope = get_outer_envelope(dc, fc, 1);
        envelope_free(outer_envelope);

        dispatch_context_free(dc);
    }
}

void tests_8_processes(int my_rank, struct futhark_context *fc) {
    printf("Testing dispatch_context_new with 8 processes. Dimensions = 8x8x8...\n");
    {
        int dimensions[3] = {8, 8, 8};
        struct dispatch_context *dc = dispatch_context_new(dimensions, MPI_INT, 3);
        chunk_info_t ci = get_chunk_info(dc);
        init_chunk_data(&ci);

        envelope_t *outer_envelope = get_outer_envelope(dc, fc, 1);

        if (my_rank == 0) {
            int expected_back_surface[] = {0, 1, 2, 3, 1, 2, 3, 4, 2, 3, 4, 5, 3, 4, 5, 6};
            int expected_front_surface[] = {3, 4, 5, 6, 4, 5, 6, 7, 5, 6, 7, 8, 6, 7, 8, 9};
            int expected_right_surface[] = {0, 1, 2, 3, 1, 2, 3, 4, 2, 3, 4, 5, 3, 4, 5, 6};
            int expected_left_surface[] = {3, 4, 5, 6, 4, 5, 6, 7, 5, 6, 7, 8, 6, 7, 8, 9};
            int expected_top_surface[] = {3,4,5,6, 4,5,6,7, 5,6,7,8, 6,7,8,9};
            int expected_bottom_surface[] = {0,1,2,3, 1,2,3,4, 2,3,4,5, 3,4,5,6};
            int expected_back_east[] = {0,1,2,3};
            int expected_back_west[] = {3,4,5,6};
            int expected_front_east[] = {3,4,5,6};
            int expected_front_west[] = {6,7,8,9};
            int expected_bottom_ne[] = {0};
            int expected_bottom_nw[] = {3};
            int expected_bottom_se[] = {3};
            int expected_bottom_sw[] = {6};
            int expected_bottom_west[] = {3,4,5,6};
            int expected_bottom_east[] = {0,1,2,3};
            int expected_bottom_north[] = {0,1,2,3};
            int expected_bottom_south[] = {3,4,5,6};
            int expected_top_ne[] = {3};
            int expected_top_nw[] = {6};
            int expected_top_se[] = {6};
            int expected_top_sw[] = {9};
            int expected_top_west[] = {6,7,8,9};
            int expected_top_east[] = {3,4,5,6};
            int expected_top_north[] = {3,4,5,6};
            int expected_top_south[] = {6,7,8,9};

            assert(memcmp(outer_envelope->back->surface->data, expected_back_surface, sizeof(expected_back_surface)) ==
                   0);
            assert(memcmp(outer_envelope->back->east->data, expected_back_east, sizeof(expected_back_east)) ==
                   0);
            assert(memcmp(outer_envelope->back->west->data, expected_back_west, sizeof(expected_back_west)) ==
                   0);

            assert(memcmp(outer_envelope->front->surface->data, expected_front_surface, sizeof(expected_front_surface)) ==
                   0);
            assert(memcmp(outer_envelope->front->east->data, expected_front_east, sizeof(expected_front_east)) ==
                   0);
            assert(memcmp(outer_envelope->front->west->data, expected_front_west, sizeof(expected_front_west)) ==
                   0);

            assert(memcmp(outer_envelope->right->surface->data, expected_right_surface, sizeof(expected_right_surface)) ==
                   0);
            assert(memcmp(outer_envelope->left->surface->data, expected_left_surface, sizeof(expected_left_surface)) ==
                   0);

            assert(memcmp(outer_envelope->bottom->surface->data, expected_bottom_surface, sizeof(expected_bottom_surface)) ==
                   0);
            assert(memcmp(outer_envelope->bottom->north_east->data, expected_bottom_ne, sizeof(expected_bottom_ne)) ==
                   0);
            assert(memcmp(outer_envelope->bottom->north_west->data, expected_bottom_nw, sizeof(expected_bottom_nw)) ==
                   0);
            assert(memcmp(outer_envelope->bottom->south_east->data, expected_bottom_se, sizeof(expected_bottom_se)) ==
                   0);
            assert(memcmp(outer_envelope->bottom->north->data, expected_bottom_north, sizeof(expected_bottom_north)) ==
                   0);
            assert(memcmp(outer_envelope->bottom->west->data, expected_bottom_west, sizeof(expected_bottom_west)) ==
                   0);
            assert(memcmp(outer_envelope->bottom->south->data, expected_bottom_south, sizeof(expected_bottom_south)) ==
                   0);
            assert(memcmp(outer_envelope->bottom->east->data, expected_bottom_east, sizeof(expected_bottom_east)) ==
                   0);

            assert(memcmp(outer_envelope->top->surface->data,    expected_top_surface, sizeof(expected_top_surface)) ==
                   0);
            assert(memcmp(outer_envelope->top->south_west->data, expected_top_sw, sizeof(expected_top_sw)) ==
                   0);
            assert(memcmp(outer_envelope->top->north_east->data, expected_top_ne, sizeof(expected_top_ne)) ==
                   0);
            assert(memcmp(outer_envelope->top->north_west->data, expected_top_nw, sizeof(expected_top_nw)) ==
                   0);
            assert(memcmp(outer_envelope->top->south_east->data, expected_top_se, sizeof(expected_top_se)) ==
                   0);
            assert(memcmp(outer_envelope->top->south_west->data, expected_top_sw, sizeof(expected_top_sw)) ==
                   0);
            assert(memcmp(outer_envelope->top->north->data, expected_top_north, sizeof(expected_top_north)) ==
                   0);
            assert(memcmp(outer_envelope->top->west->data, expected_top_west, sizeof(expected_top_west)) ==
                   0);
            assert(memcmp(outer_envelope->top->south->data, expected_top_south, sizeof(expected_top_south)) ==
                   0);
            assert(memcmp(outer_envelope->top->east->data, expected_top_east, sizeof(expected_top_east)) ==
                   0);

        }

        envelope_free(outer_envelope);
        dispatch_context_free(dc);
    }
}

int main(int argc, char *argv[]) {
    MPI_Init(&argc, &argv);
    struct futhark_context_config *fcc = futhark_context_config_new();
    struct futhark_context *fc = futhark_context_new(fcc);

    int my_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

    int world_size;
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    switch (world_size) {
        case 1:
            tests_1_process(fc);
            break;
        case 2:
            tests_2_processes(my_rank, fc);
            break;
        case 3:
            tests_3_processes(my_rank, fc);
            break;
        case 4:
            tests_4_processes(my_rank, fc);
            break;
        case 6:
            tests_6_processes(my_rank, fc);
            break;
        case 8:
            tests_8_processes(my_rank, fc);
            break;
        default:
            break;
    }

    futhark_context_config_free(fcc);
    futhark_context_free(fc);
    MPI_Finalize();
}