#include <mpi.h>

#include <algorithm>
#include <cmath>
#include <iostream>
#include <memory>

#define ALPHA 100000.f              // параметр правой чачти
#define N 144                       // размерность сетки (количество ячеек)
#define EPS 1E-8                    // условие отстановы итерационного метода
#define MAX_ITERS 10000             // максимальное число итераций

#define EPS_COMP 1E-2               // погрешность сравнения результатов

#define MASTER 0                    // ID мастер-процесса

// глобальные переменные
const double h = 2.0 / (N - 1.0);   // размер шага вдоль осей x, y и z
int i0;                             // переменная для сопоставления отдельных блоков цельному параллелепипеду
int current_iteration;              // текущая итерация
int my_rank, n_procs;               // данный процесс, число процессов

// искомая функция фи
double CalculateFunc(const double x, const double y, const double z);
// правая часть уравнения
double CalculateRo(const double x, const double y, const double z);
// инициализация 3-хмерного массива + инициализация границ 3х-мерного пространства
void InitBounds(double* A, double* f, const int Im_min, const int Im_max, const int Im_bound_up, const int Im_bound_down);
// перегрузка, что не пересчитывать f несколько раз
void InitBounds(double* A, const int Im_min, const int Im_max, const int Im_bound_up, const int Im_bound_down);
// рассчет новых элементов матрицы A перед обменом
double DoYakobiIteration(const int imin, const int imax, double* A, double* A_next, double* f, const double _covergence_local);
// непараллельное решение задачи методом Якоби
double CalculateSequentialJacobi();
// параллельное решение задачи методом Якоби
double CalculateParallelJacobi();
// проверка полученного результата
void CompareResults(double* result, const int step_max);

int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);                     // инициализация MPI
    MPI_Comm_size(MPI_COMM_WORLD, &n_procs);    // определение числа процессов в области связи
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);    // определение номера вызывающего процесса
    if (my_rank == MASTER) std::cout << "Всего процессов: " << n_procs << std::endl << std::endl;

    double covergence;
    double t_start, t_end;
    if (my_rank == MASTER && n_procs == 1) {
        MPI_Barrier(MPI_COMM_WORLD);
        t_start = MPI_Wtime();
        covergence = CalculateSequentialJacobi();
        MPI_Barrier(MPI_COMM_WORLD);
        t_end = MPI_Wtime();
        std::cout << "Условие сходимости: " << covergence << std::endl;
        std::cout << "Число итераций: " << current_iteration << std::endl;
        std::cout << "Затраченное время: " << t_end - t_start << std::endl;
    } else if (N % n_procs == 0 && n_procs > 1) {
        MPI_Barrier(MPI_COMM_WORLD);
        t_start = MPI_Wtime();
        covergence = CalculateParallelJacobi();
        MPI_Barrier(MPI_COMM_WORLD);
        t_end = MPI_Wtime();
        if (my_rank == MASTER) {
            std::cout << "Условие сходимости: " << covergence << std::endl;
            std::cout << "Число итераций: " << current_iteration << std::endl;
            std::cout << "Затраченное время: " << t_end - t_start << std::endl;
        }
    } else {
        if (my_rank == MASTER)
            std::cout << "Программа завершена, скорректруйте N (N % n_procs == 0)!" << std::endl;
        MPI_Finalize();
        return -1;
    }

prog_exit:
    MPI_Finalize();
    return 0;
}

// искомая функция фи
double CalculateFunc(const double x, const double y, const double z) {
    return x * x + y * y + z * z;
};

// правая часть уравнения
double CalculateRo(const double x, const double y, const double z) {
    return 6 - ALPHA * (x * x + y * y + z * z);
}

// инициализация 3-хмерного массива + инициализация границ 3х-мерного пространства
void InitBounds(double* A, double* f, const int Im_min, const int Im_max, const int Im_bound_up, const int Im_bound_down) {
    for (int i = Im_min; i < Im_max; i++)
        for (int j = 0; j < N; j++)
            for (int k = 0; k < N; k++) {
                // на границе
                if (i == Im_bound_up || j == 0 || k == 0 || i == Im_bound_down || j == N - 1 || k == N - 1)
                    A[i * N * N + j * N + k] = CalculateFunc((i + i0) * h, j * h, k * h);
                else
                    // внутри области решения
                    A[i * N * N + j * N + k] = 0.0;
                // рассчитываем правые части
                f[i * N * N + j * N + k] = CalculateRo((i + i0) * h, j * h, k * h);
            }
};

// перегрузка функции
void InitBounds(double* A, const int Im_min, const int Im_max, const int Im_bound_up, const int Im_bound_down) {
    for (int i = Im_min; i < Im_max; i++)
        for (int j = 0; j < N; j++)
            for (int k = 0; k < N; k++) {
                if (i == Im_bound_up || j == 0 || k == 0 || i == Im_bound_down || j == N - 1 || k == N - 1)
                    A[i * N * N + j * N + k] = CalculateFunc((i + i0) * h, j * h, k * h);
                else
                    A[i * N * N + j * N + k] = 0.0;
            }
};

// рассчет новых элементов матрицы A перед обменом
double DoYakobiIteration(const int imin, const int imax, double* A, double* A_next, double* f, const double _covergence_local) {
    double covergence_local = _covergence_local;
    double coef = 1.0 / ((6.0 / h * h) + ALPHA);
    double F1, F2, F3;
    for (int i = imin; i < imax; i++) {
        for (int j = 1; j < N - 1; j++) {
            for (int k = 1; k < N - 1; k++) {
                F1 = A[(i - 1) * N * N + j * N + k] + A[(i + 1) * N * N + j * N + k];
                F2 = A[i * N * N + (j - 1) * N + k] + A[i * N * N + (j + 1) * N + k];
                F3 = A[i * N * N + j * N + k - 1] + A[i * N * N + j * N + k + 1];
                A_next[i * N * N + j * N + k] = coef * (F1 + F2 + F3 - f[i * N * N + j * N + k]);

                covergence_local = std::max(covergence_local, fabs(A_next[i * N * N + j * N + k] - A[i * N * N + j * N + k]));
            }
        }
    }
    return covergence_local;
};

// непараллельное решение задачи методом Якоби
double CalculateSequentialJacobi() {
    double* A = new double[N * N * N];       // старые значения
    double* A_next = new double[N * N * N];  // новые значения
    double* f = new double[N * N * N];       // рассчитанные правые части
    double covergence, covergence_local;

    i0 = 0;
    int Im_min = 0;
    int bound_up = 0;
    int bound_down = N - 1;

    InitBounds(A, f, Im_min, N, bound_up, bound_down);
    InitBounds(A_next, Im_min, N, bound_up, bound_down);

    current_iteration = 0;
    covergence_local = 0.f;
    // итерации метода Якоби
    do {
        covergence_local = 0;
        covergence_local = DoYakobiIteration(1, N - 1, A, A_next, f, covergence_local);
        std::swap(A, A_next);                   // обмен значениями
        covergence = covergence_local;
        current_iteration++;
    } while (covergence > EPS && current_iteration < MAX_ITERS);

    delete[] A_next;
    delete[] f;
    // return A;
    return covergence;
};

// параллельное решение задачи методом Якоби
double CalculateParallelJacobi() {
    // идентификаторы для перессылки сообщений
    MPI_Request up_s, down_s;
    MPI_Request up_r, down_r;

    int Im_min, Im_max;
    int bound_up, bound_down;
    double covergence;              // условие выхода
    int Im_block;                   // размер блока для процесса

    // место под выделаемые из общего параллелепипеда блоки
    double* A;                      // старые значения
    double* A_next;                 // новые значения

    double* f;                      // рассчитанные правые части

    // выделение определённому процессу куска
    // граничный блок
    if (my_rank == 0 || my_rank == n_procs - 1)
        Im_block = N / n_procs + 1;
    else
        //внутренний блок
        Im_block = N / n_procs + 2;

    A = new double[Im_block * N * N];
    A_next = new double[Im_block * N * N];
    f = new double[Im_block * N * N];
    i0 = my_rank * N / n_procs - (my_rank != 0);

    if (my_rank == 0) {
        Im_min = 0;
        Im_max = Im_block;
        bound_up = 0;
        bound_down = -1;
    } else if (my_rank == n_procs - 1) {
        Im_min = 0;
        Im_max = Im_block;
        bound_up = -1;
        bound_down = Im_block - 1;
    } else {
        Im_min = 0;
        Im_max = Im_block;
        bound_up = -1;
        bound_down = -1;
    }

    // здесь вычисляются границы для каждого блока
    // для потоков больше 0 границы вычисляются неверно,
    // поэтому прежде чем вычислять внутренние точки методом Якоби,
    // нужно принять от других потоков внешние границы
    InitBounds(A, f, Im_min, Im_max, bound_up, bound_down);
    InitBounds(A_next, Im_min, Im_max, bound_up, bound_down);

    current_iteration = 0;
    do {
        current_iteration++;
        double covergence_local = 0.f;
        int _i = 1;
        int _n = Im_block - 1;
        if (my_rank != 0 && N / n_procs != 1) {
            int i = 1;
            _i++;
            covergence_local = DoYakobiIteration(i, i + 1, A, A_next, f, covergence_local);
            MPI_Isend(&A_next[N * N], N * N, MPI_DOUBLE, my_rank - 1, 0, MPI_COMM_WORLD, &down_s);
            MPI_Irecv(&A_next[0], N * N, MPI_DOUBLE, my_rank - 1, 0, MPI_COMM_WORLD, &down_r);
        }

        if (my_rank != n_procs - 1 && N / n_procs != 1) {
            int i = Im_block - 2;
            _n--;
            covergence_local = DoYakobiIteration(i, i + 1, A, A_next, f, covergence_local);
            // неблокирующая передача сообщения
            MPI_Isend(&A_next[(Im_block - 2) * N * N], N * N, MPI_DOUBLE, my_rank + 1, 0, MPI_COMM_WORLD, &up_s);
            // неблокирующий приём сообщения
            MPI_Irecv(&A_next[(Im_block - 1) * N * N], N * N, MPI_DOUBLE, my_rank + 1, 0, MPI_COMM_WORLD, &up_r);
        }
        covergence_local = DoYakobiIteration(_i, _n, A, A_next, f, covergence_local);
        // для (N/n_procs == 1) нужно сначала, наверное, принимать сообщения о внешних границах,
        // потом вызвать метод DoYakobiIteration,
        // потом отправлять вычисленную (внутреннюю) границу
        if (N / n_procs == 1) {
            if (my_rank == 0)
                _i = 0;
            if (my_rank != 0) {
                // неблокирующая передача сообщения
                MPI_Isend(&A_next[_i * N * N], N * N, MPI_DOUBLE, my_rank - 1, 0, MPI_COMM_WORLD, &down_s);
                // неблокирующий приём сообщения
                MPI_Irecv(&A_next[(_i - 1) * N * N], N * N, MPI_DOUBLE, my_rank - 1, 0, MPI_COMM_WORLD, &down_r);
            }
            if (my_rank != n_procs - 1) {
                // неблокирующая передача сообщения
                // N * N – количество точек на границе (YxZ),
                // эта граница образуется при пересечении параллелепипеда с плоскостью х = _i.
                MPI_Isend(&A_next[_i * N * N], N * N, MPI_DOUBLE, my_rank + 1, 0, MPI_COMM_WORLD, &up_s);
                // неблокирующий приём сообщения
                MPI_Irecv(&A_next[(_i + 1) * N * N], N * N, MPI_DOUBLE, my_rank + 1, 0, MPI_COMM_WORLD, &up_r);
            }
        }
        // cохранение результата в адресном пространстве всех процессов (covergence_local -> covergence)
        MPI_Allreduce(&covergence_local, &covergence, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);

        // если действие производится не над первым процессом
        if (my_rank != 0) {
            MPI_Wait(&down_s, MPI_STATUS_IGNORE);
            MPI_Wait(&down_r, MPI_STATUS_IGNORE);
        }
        // eсли действие производится не над последним процессом
        if (my_rank != n_procs - 1) {
            MPI_Wait(&up_s, MPI_STATUS_IGNORE);
            MPI_Wait(&up_r, MPI_STATUS_IGNORE);
        }
        std::swap(A, A_next);               // обмен значениями
        covergence = covergence_local;
    } while (covergence >= EPS && current_iteration < MAX_ITERS);

    // проверка результата
    if (my_rank == 0)
        Im_min = 0;
    else {
        if (my_rank == n_procs - 1)
            Im_min = 1;
        else {
            Im_min = 1;
            Im_max--;
        }
    }
    CompareResults(A, Im_max);
    delete[] A_next;
    delete[] f;

    // return A;
    return covergence;
};

// проверка результатов
void CompareResults(double* result, const int step_max) {
    const int step_min = 0;
    int error_count = 0;
    for (int i = i0 + step_min, p = step_min; i < N && p < step_max; i++, p++) {
        double x = i * h;
        for (int j = 0; j < N; j++) {
            double y = j * h;
            for (int k = 0; k < h; k++) {
                double z = k * h;
                if (fabs(result[p * N * N + j * N + k] - CalculateFunc(x, y, z)) >= EPS_COMP) {
                    std::cout << i * N * N + j * N + k << "\t";
                    std::cout << CalculateFunc(x, y, z) << "\t\t";
                    std::cout << p * N * N + j * N + k << "\t";
                    std::cout << result[p * N * N + j * N + k] << "\t";
                    std::cout << "Error!" << std::endl;
                    error_count++;
                }
            }
        }
    }
    if (error_count != 0)
        std::cout << error_count << " Ошибка при работе процесса " << my_rank << "!" << std::endl;
};