/* Include benchmark-specific header. */
#include "heat-3d.h"
#include <unistd.h>
#include <mpi.h>

double bench_t_start, bench_t_end;

static
double rtclock()
{
    struct timeval Tp;
    int stat;
    stat = gettimeofday(&Tp, NULL);
    if (stat != 0) {
        printf("Error return from gettimeofday: %d", stat);
    }
    return (Tp.tv_sec + Tp.tv_usec * 1.0e-6);
}

void
bench_timer_start()
{
    bench_t_start = rtclock();
}

void
bench_timer_stop()
{
    bench_t_end = rtclock();
}

void
old_bench_timer_print()
{
    printf("Time in seconds = %0.6lf\n", bench_t_end - bench_t_start);
}

void
bench_timer_print()
{
    printf("%0.6lf\n", bench_t_end - bench_t_start);
}

static
void init_subarray(int n, int n1, int n2, int n3,
        double A[n1][n2][n3], double B[n1][n2][n3],
        int i_offset, int j_offset, int k_offset,
        int i_stop, int j_stop, int k_stop)
{
    i_offset = (i_offset < 0) ? 0 : i_offset;
    j_offset = (j_offset < 0) ? 0 : j_offset;
    k_offset = (k_offset < 0) ? 0 : k_offset;

    for (int i = i_offset; i < i_stop; i++) {
        for (int j = j_offset; j < j_stop; j++) {
            for (int k = k_offset; k < k_stop; k++) {
                A[i][j][k] = B[i][j][k] = (double)(i + j + (n - k)) * 10 / (n);
            }
        }
    }
}

static
void print_array(int n,
        double A[n][n][n])
{
    fprintf(stderr, "==BEGIN DUMP_ARRAYS==\n");
    fprintf(stderr, "begin dump: %s", "A");
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            for (int k = 0; k < n; k++) {
                if ((i * n * n + j * n + k) % 20 == 0) {
                    fprintf(stderr, "\n");
                }
                fprintf(stderr, "%0.2lf ", A[i][j][k]);
            }
        }
    }
    fprintf(stderr, "\nend   dump: %s\n", "A");
    fprintf(stderr, "==END   DUMP_ARRAYS==\n");
}

static
void kernel_heat_3d(int tsteps,
        int n,
        double A[n][n][n],
        double B[n][n][n])
{
    for (int t = 1; t <= TSTEPS; t++) {
        for (int i = 1; i < n - 1; i++) {
            for (int j = 1; j < n - 1; j++) {
                for (int k = 1; k < n - 1; k++) {
                    B[i][j][k] = 0.125 * (A[i + 1][j][k] - 2.0 * A[i][j][k] + A[i - 1][j][k])
                            + 0.125 * (A[i][j + 1][k] - 2.0 * A[i][j][k] + A[i][j - 1][k])
                            + 0.125 * (A[i][j][k + 1] - 2.0 * A[i][j][k] + A[i][j][k - 1])
                            + A[i][j][k];
                }
            }
        }
        for (int i = 1; i < n - 1; i++) {
            for (int j = 1; j < n - 1; j++) {
                for (int k = 1; k < n - 1; k++) {
                    A[i][j][k] = 0.125 * (B[i + 1][j][k] - 2.0 * B[i][j][k] + B[i - 1][j][k])
                            + 0.125 * (B[i][j + 1][k] - 2.0 * B[i][j][k] + B[i][j - 1][k])
                            + 0.125 * (B[i][j][k + 1] - 2.0 * B[i][j][k] + B[i][j][k - 1])
                            + B[i][j][k];
                }
            }
        }
    }
}


void
test(int data_size, int* argc_address, char*** argv_address)
{
    int n = data_size;
    int tsteps = TSTEPS;

    MPI_Init(argc_address, argv_address);

    const int ndims = 3;
    int rank, procs;
    MPI_Comm cart_comm;
    int reorder;
    int coord[ndims], id;
    int right, down, left, up, next, previous;

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &procs);

    /*
    Мы хотим разделить наш 3-х мерный тензор на кубики, которые мы будем считать
    на каждом из процессоров. Сколько будем кубиков (в зависимости от числа
    процессов) по каждой из размерностей можно посчитать с помошью MPI_Dims_create,
    хотя вручную тоже можно было это написать.
    */

    int dim[/*ndims*/] = { 0, 0, 0 };
    MPI_Dims_create(procs, ndims, dim);

    /*
    Периодичность по какой-то из размерностей означает, что следующим после
    последнего будет первый элемент. У нас ни по одной размерностей периодичностей
    нет.
    */

    int periodical[/*ndims*/] = { 0, 0, 0 };
    reorder = 0;

    MPI_Cart_create(MPI_COMM_WORLD, ndims, dim, periodical, reorder, &cart_comm);
    MPI_Cart_coords(cart_comm, rank, ndims, coord);

    // Удобный доступ к рангу соседей по каждому из направлений:
    MPI_Cart_shift(cart_comm, 0, 1, &up, &down);
    MPI_Cart_shift(cart_comm, 1, 1, &left, &right);
    MPI_Cart_shift(cart_comm, 2, 1, &previous, &next);

    // + 2 нужно для теневых граней
    int n1 = n / dim[0] + 2, n2 = n / dim[1] + 2, n3 = n / dim[2] + 2;

    // Не всегда n % dim[i] == 0 (i = 0, 1, 2)
    if (down == -1) {
        n1 += n % dim[0];
    }
    if (right == -1) {
        n2 += n % dim[1];
    }
    if (next == -1) {
        n3 += n % dim[2];
    }

    /*
    Теперь надо создать на каждом из MPI-процессов A и B и их проинициализировать
    */
    double (*A)[n1][n2][n3]; A = (double(*)[n1][n2][n3])malloc((n1) * (n2) * (n3) * sizeof(double));
    double (*B)[n1][n2][n3]; B = (double(*)[n1][n2][n3])malloc((n1) * (n2) * (n3) * sizeof(double));

    // Вычитаем -1 для теневых граней
    int i_offset = n / dim[0] * coord[0] - 1;
    int j_offset = n / dim[1] * coord[1] - 1;
    int k_offset = n / dim[2] * coord[2] - 1;

    int i_bound = i_offset + n1;
    int j_bound = j_offset + n2;
    int k_bound = k_offset + n3;

    for (int i = i_offset; i < i_bound; i++) {
        for (int j = j_offset; j < j_bound; j++) {
            for (int k = k_offset; k < k_bound; k++) {
                (*A)[i - i_offset][j - j_offset][k - k_offset]
                        = (*B)[i - i_offset][j - j_offset][k - k_offset]
                        = (double)(i + j + (n - k)) * 10 / (n);
            }
        }
    }

    //bench_timer_start();

    MPI_Request requests[12]; // У нас 6 shadow edges, один request на send, другой на recieve
    MPI_Status statuses[12];

    /*
    Нам придётся передавать грани, то есть 2D подмассивы. Для этого нам понадобится сразу 3 типа.
    Так как элементы разных теневых граней, перпендикулярных разным осям по-разному лежат в памяти.
    */

    MPI_Datatype matrix_oxy, matrix_oxz, matrix_oyz;

    int sizes[/*ndims*/] = { n1, n2, n3 };
    int subsizes_oxy[/*ndims*/] = { n1, n2, 1 };
    int subsizes_oxz[/*ndims*/] = { n1, 1, n3 };
    int subsizes_oyz[/*ndims*/] = { 1, n2, n3 };
    int starts[/*ndims*/] = { 0, 0, 0 };

    MPI_Type_create_subarray(ndims, sizes, subsizes_oxy, starts, MPI_ORDER_C, MPI_DOUBLE, &matrix_oxy);
    MPI_Type_commit(&matrix_oxy);

    MPI_Type_create_subarray(ndims, sizes, subsizes_oxz, starts, MPI_ORDER_C, MPI_DOUBLE, &matrix_oxz);
    MPI_Type_commit(&matrix_oxz);

    MPI_Type_create_subarray(ndims, sizes, subsizes_oyz, starts, MPI_ORDER_C, MPI_DOUBLE, &matrix_oyz);
    MPI_Type_commit(&matrix_oyz);

    double time_start = MPI_Wtime();

    for (int t = 1; t <= TSTEPS; t++) {
        for (int i = 1 + (up == -1); i < n1 - 1 - (down == -1); i++) {
            for (int j = 1 + (left == -1); j < n2 - 1 - (right == -1); j++) {
                for (int k = 1 + (previous == -1); k < n3 - 1 - (next == -1); k++) {
                    (*B)[i][j][k] = 0.125 * ((*A)[i + 1][j][k] - 2.0 * (*A)[i][j][k] + (*A)[i - 1][j][k])
                            + 0.125 * ((*A)[i][j + 1][k] - 2.0 * (*A)[i][j][k] + (*A)[i][j - 1][k])
                            + 0.125 * ((*A)[i][j][k + 1] - 2.0 * (*A)[i][j][k] + (*A)[i][j][k - 1])
                            + (*A)[i][j][k];
                }
            }
        }
        /*
        У нас поменялся массив B. Соседи должны обновить значения теневых граней.
        Также нужно подождать, пока обмен с соседями не завершится. Ведь для дальнейших вычислений
        нам понадабятся данные, полученные с других процессов
        */

        // Получаем данные от верхнего соседа и передаём нижнему (по 1-му измерению)
        MPI_Irecv(&(*B)[0][0][0], 1, matrix_oyz, up, 1215, MPI_COMM_WORLD, &requests[0]);
        MPI_Isend(&(*B)[n1 - 2][0][0], 1, matrix_oyz, down, 1215, MPI_COMM_WORLD, &requests[1]);

        // Получаем данные от нижнего соседа и передаём верхнему (по 1-му измерению)
        MPI_Irecv(&(*B)[n1 - 1][0][0], 1, matrix_oyz, down, 1216, MPI_COMM_WORLD, &requests[2]);
        MPI_Isend(&(*B)[1][0][0], 1, matrix_oyz, up, 1216, MPI_COMM_WORLD, &requests[3]);

        // Получаем данные от правого соседа и передаём левому (по 2-му измерению)
        MPI_Irecv(&(*B)[0][0][0], 1, matrix_oxz, right, 1217, MPI_COMM_WORLD, &requests[4]);
        MPI_Isend(&(*B)[0][n2 - 2][0], 1, matrix_oxz, left, 1217, MPI_COMM_WORLD, &requests[5]);

        // Получаем данные от левого соседа и передаём правому (по 2-му измерению)
        MPI_Irecv(&(*B)[0][n2 - 1][0], 1, matrix_oxz, left, 1218, MPI_COMM_WORLD, &requests[6]);
        MPI_Isend(&(*B)[0][1][0], 1, matrix_oxz, right, 1218, MPI_COMM_WORLD, &requests[7]);

        // Получаем данные от следующего соседа и передаём предыдущему (по 3-му измерению)
        MPI_Irecv(&(*B)[0][0][0], 1, matrix_oxy, next, 1219, MPI_COMM_WORLD, &requests[8]);
        MPI_Isend(&(*B)[0][0][n3 - 2], 1, matrix_oxy, previous, 1219, MPI_COMM_WORLD, &requests[9]);

        // Получаем данные от предыдущего соседа и передаём следующему (по 3-му измерению)
        MPI_Irecv(&(*B)[0][0][n3 - 1], 1, matrix_oxy, previous, 1220, MPI_COMM_WORLD, &requests[10]);
        MPI_Isend(&(*B)[0][0][1], 1, matrix_oxy, next, 1220, MPI_COMM_WORLD, &requests[11]);

        // Дожидаемся получения всех теневых граней
        MPI_Waitall(12, requests, statuses);

        for (int i = 1 + (up == -1); i < n1 - 1 - (down == -1); i++) {
            for (int j = 1 + (left == -1); j < n2 - 1 - (right == -1); j++) {
                for (int k = 1 + (previous == -1); k < n3 - 1 - (next == -1); k++) {
                    (*A)[i][j][k] = 0.125 * ((*B)[i + 1][j][k] - 2.0 * (*B)[i][j][k] + (*B)[i - 1][j][k])
                            + 0.125 * ((*B)[i][j + 1][k] - 2.0 * (*B)[i][j][k] + (*B)[i][j - 1][k])
                            + 0.125 * ((*B)[i][j][k + 1] - 2.0 * (*B)[i][j][k] + (*B)[i][j][k - 1])
                            + (*B)[i][j][k];
                }
            }
        }

        /*
        Теперь поменялся массив A. Нужно обновить значение теневых граней.
        */

        MPI_Irecv(&(*A)[0][0][0], 1, matrix_oyz, up, 1221, MPI_COMM_WORLD, &requests[0]);
        MPI_Isend(&(*A)[n1 - 2][0][0], 1, matrix_oyz, down, 1221, MPI_COMM_WORLD, &requests[1]);

        // Получаем данные от нижнего соседа и передаём верхнему (по 1-му измерению)
        MPI_Irecv(&(*A)[n1 - 1][0][0], 1, matrix_oyz, down, 1222, MPI_COMM_WORLD, &requests[2]);
        MPI_Isend(&(*A)[1][0][0], 1, matrix_oyz, up, 1222, MPI_COMM_WORLD, &requests[3]);

        // Получаем данные от правого соседа и передаём левому (по 2-му измерению)
        MPI_Irecv(&(*A)[0][0][0], 1, matrix_oxz, right, 1223, MPI_COMM_WORLD, &requests[4]);
        MPI_Isend(&(*A)[0][n2 - 2][0], 1, matrix_oxz, left, 1223, MPI_COMM_WORLD, &requests[5]);

        // Получаем данные от левого соседа и передаём правому (по 2-му измерению)
        MPI_Irecv(&(*A)[0][n2 - 1][0], 1, matrix_oxz, left, 1224, MPI_COMM_WORLD, &requests[6]);
        MPI_Isend(&(*A)[0][1][0], 1, matrix_oxz, right, 1224, MPI_COMM_WORLD, &requests[7]);

        // Получаем данные от следующего соседа и передаём предыдущему (по 3-му измерению)
        MPI_Irecv(&(*A)[0][0][0], 1, matrix_oxy, next, 1225, MPI_COMM_WORLD, &requests[8]);
        MPI_Isend(&(*A)[0][0][n3 - 2], 1, matrix_oxy, previous, 1225, MPI_COMM_WORLD, &requests[9]);

        // Получаем данные от предыдущего соседа и передаём следующему (по 3-му измерению)
        MPI_Irecv(&(*A)[0][0][n3 - 1], 1, matrix_oxy, previous, 1226, MPI_COMM_WORLD, &requests[10]);
        MPI_Isend(&(*A)[0][0][1], 1, matrix_oxy, next, 1226, MPI_COMM_WORLD, &requests[11]);

        // Дожидаемся получения всех теневых граней
        MPI_Waitall(12, requests, statuses);
    }

    double time_end = MPI_Wtime();
    if (rank == 0) {
        printf("n, procs, time = %d, %d, %lf\n", n, procs, time_end - time_start);
    }

    /*
    if (argc > 42 && !strcmp(argv[0], "")) {
        print_array(n, *A);
    }
    */

    free((void*)A);
    free((void*)B);

    MPI_Finalize();
}

int
main(int argc, char** argv)
{
    int data_size[] = { 100 };
    int data_sizes_length = sizeof(data_size) / sizeof(*data_size);

    for (int j = 0; j < data_sizes_length; j++) {
        printf("%d\n", data_size[j]);
        test(data_size[j], &argc, &argv);
    }
    return 0;
}