#include <iostream>
#include <vector>
#include <cmath>
#include <chrono>
#include <omp.h>
#include <cstring>
#include <mpi.h>
#include <fstream>

#define EPS 1E-7
#define PI 3.14159265358979323846

using namespace std;

bool CheckSphereIn(const double& R, const double& x, const double& y){
    return x*x + y*y - R*R <= EPS;
}

int MakeSpherePointsField(const double& R, const double& h, vector<double>& X, vector<double>& Y,
                        const double& start_x=0.f, const double& start_y=0.f){ // 0.f - это преобразование 0 к double(чтобы было 0.0)
    if (abs(R) < EPS)
        return EXIT_FAILURE;
    if (abs(h) < EPS)
        return EXIT_FAILURE;

    const auto n = static_cast<long>(R / h);

    X.clear();
    Y.clear();

    for (auto i = 0; i < n; ++i)
        for (auto j = 0; j < n; ++j)
            if (CheckSphereIn(R, h*i, h*j)){
                X.emplace_back(h*i + start_x);
                Y.emplace_back(h*j + start_y);
            }

    return EXIT_SUCCESS;
}

double RKF(const double& y, const double& f, const double& h){
    return y + h*f;
}

int set_tasks(const int n, const int k, vector<int>& left, vector<int>& right){
    left.clear();
    right.clear();

    for (int i = 0; i < k; ++i){
        left.push_back((int) (((long)i * n) / k));
        right.push_back((int) (((long) (i + 1) * n) / k - 1));
    }
    return 0;
}

vector<int> left_indexes, right_indexes;

// добавить омега
int ToNextStepMPI(double* X, double* Y, const unsigned long n, const double& h, const int& id, const int& comm_size){
    if (abs(h) < EPS)
        return EXIT_FAILURE;

    double *oldX = new double[n], *oldY = new double[n];
    memcpy(oldX, X, sizeof(double)*n);
    memcpy(oldY, Y, sizeof(double)*n);

#pragma omp parallel for
    for (auto i = left_indexes[id]; i < right_indexes[id]; ++i){
    	double k1_x, k2_x, k3_x, k4_x;
    	double k1_y, k2_y, k3_y, k4_y;
        double sum_x = 0, sum_y = 0;
        for (auto j = 0; j < n; ++j){
            if (i == j)
                continue;
            double d = pow(oldX[i] - oldX[j], 2) + pow(oldY[i] - oldY[j], 2); //расстояние между двумя точками
            sum_x += (oldY[i] - oldY[j]) / d;
            sum_y += (oldX[i] - oldX[j]) / d;
        }

        k1_x = sum_x/-2/PI;
        k1_y = sum_y/2/PI;

        double c_x = k1_x*h/2;
        double c_y = k1_y*h/2;
        for (auto j = 0; j < n; ++j){
            if (i == j)
                continue;
            double d = pow(oldX[i]+c_x - oldX[j]+c_x, 2) + pow(oldY[i]+c_y - oldY[j]+c_y, 2); //расстояние между двумя точками
            sum_x += (oldY[i]+c_y - oldY[j]+c_y) / d;
            sum_y += (oldX[i]+c_x - oldX[j]+c_x) / d;
        }
        k2_x = sum_x/-2/PI;
        k2_y = sum_y/-2/PI;

        c_x = k2_x*h/2;
        c_y = k2_y*h/2;
        for (auto j = 0; j < n; ++j){
            if (i == j)
                continue;
            double d = pow(oldX[i]+c_x - oldX[j]+c_x, 2) + pow(oldY[i]+c_y - oldY[j]+c_y, 2); //расстояние между двумя точками
            sum_x += (oldY[i]+c_y - oldY[j]+c_y) / d;
            sum_y += (oldX[i]+c_x - oldX[j]+c_x) / d;
        }
     	k3_x = sum_x/-2/PI;
        k3_y = sum_y/2/PI;

        c_x = k3_x*h;
        c_y = k3_y*h;
        for (auto j = 0; j < n; ++j){
            if (i == j)
                continue;
            double d = pow(oldX[i] - oldX[j], 2) + pow(oldY[i] - oldY[j], 2); //расстояние между двумя точками
            sum_x += (oldY[i] - oldY[j]) / d;
            sum_y += (oldX[i] - oldX[j]) / d;
        }
        k4_x = sum_x/-2/PI;
        k4_x = sum_y/2/PI;

        X[i] = X[i] + (k1_x + 2*k2_x + 2*k3_x + k4_x)*h/6;
        Y[i] = Y[i] + (k1_y+ 2*k2_y + 2*k3_y + k4_y)*h/6;
    }

    // всеобщее обновление массивов X и Y
    for (auto i = 0; i < comm_size; ++i){
        MPI_Bcast(X + left_indexes[i], right_indexes[i] - left_indexes[i], MPI_DOUBLE, i, MPI_COMM_WORLD);
        MPI_Bcast(Y + left_indexes[i], right_indexes[i] - left_indexes[i], MPI_DOUBLE, i, MPI_COMM_WORLD);
    }

    delete[] oldX;
    delete[] oldY;

    return EXIT_SUCCESS;
}

int get_default_max_threads_num(){
    int res = 0;
#pragma omp parallel
    {
#pragma omp single
        res = omp_get_num_threads();
    }
    return res;
}

void check_err(string err_str, const int err){
    if (err)
        cout << err_str << " " << err << endl;
}

int main(int argc, char** argv) {


    check_err("MPI_Init", MPI_Init(&argc, &argv));
    int id, cores;

    MPI_Comm_rank(MPI_COMM_WORLD, &id);
    MPI_Comm_size(MPI_COMM_WORLD, &cores);


    vector<double> Xv, Yv;
    MakeSpherePointsField(1, 0.2, Xv, Yv); // вычисление точек, которые попадают в сферу с радиусом 5; шаг 0.05; Xv, Yv координаты этих точек

    // перенос данных из векторов в обычные массивы, чтобы вставлять часть массива в другой была O(1)
    double *X = new double[Xv.size()], *Y = new double[Yv.size()];
    int n = Xv.size();
    cout << "n?: " << n << endl;

    memcpy(X, Xv.data(), sizeof(double)*n);
    Xv.clear();
    memcpy(Y, Yv.data(), sizeof(double)*n);
    Yv.clear();

    if (id == 0) {
        cout << "Cores count: " << cores << endl;
        cout << "Threads count: " << get_default_max_threads_num() << endl;
        cout << "Points count: " << n << endl;
    }

    // создает два массива, в которых указаны индексы начала и конца точек которые будут обрабатываться в каждом узле
    set_tasks(static_cast<const int>(n), cores, left_indexes, right_indexes);

    const unsigned long steps = 100;
    const double h = 0.1;
    chrono::duration<double> timer;


    FILE* file = fopen("points.txt", "w");
    fprintf(file, "%d\n", n);
    int i;

    auto start_time = chrono::steady_clock::now();
    for (auto i = 0; i < steps; ++i)
    {
        if (id == 0) {
        	for (i = 0; i < n; ++i)
    	        	fprintf(file, "%f %f\n", X[i], Y[i]);

        }
        ToNextStepMPI(X, Y, n, h, id, cores);
    }
    fclose(file);
    timer = chrono::steady_clock::now() - start_time;


    MPI_Barrier(MPI_COMM_WORLD);

    if (id == 0)
        cout << "Time: " << timer.count() << " sec" << endl;

    delete[] X;
    delete[] Y;

    MPI_Finalize();

    return EXIT_SUCCESS;
}