#include <iostream>
#include <cmath>
#include <stdlib.h>
#include <time.h>
#include <stdio.h>
#include <mpi.h>

void Matrix_by_vector(int N, double *M, double *V, double *R)
//N - размерность, M - матрица, V - вектор, R - реузультат
{
    int rank, size;
    MPI_Status status;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    int mod = N%size;

    double *localresult = new double[N/size];
    double matrix[N][N];//local matirx
    MPI_Barrier(MPI_COMM_WORLD);//Пока все процессы не вызовут эту функцию, вызвавшие будут ждать

//Передача частей матрицы процессам(осуществляется 0-ым процессом)
    if(rank==0){
        for(int i = 0; i < N/size + ((0 < mod) ? (1) : (0)); i++)
           for(int j = 0; j < N/size; j++)
                matrix[i][j] = M[j + (N*i)];
        for(int i = 1; i < size; i++){
           MPI_Send(M+((N*N/size)*(i)), N*(N/size + ((i < mod) ? (1) : (0))), MPI_DOUBLE, i, 0, MPI_COMM_WORLD);
        }
    }
    else{
        MPI_Recv(matrix, N*(N/size + ((rank < mod) ? (1) : (0))), MPI_DOUBLE, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }

//Передача 0-ым процессом вектора, на который умножается матрица
    if(rank == 0){
        for(int i = 0; i < size; i++){
            if(i != rank){
                MPI_Send(V, N, MPI_DOUBLE, i, 0, MPI_COMM_WORLD);
            }
        }
    } else{
        MPI_Recv(V, N, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }

//Выполнение умножения матрицы на вектор
    for(int i = 0; i < N/size + ((rank < mod) ? (1) : (0)); i++){
        localresult[i] = 0;
        for(int j = 0; j < N; j++)
        {
            localresult[i] += matrix[i][j]*V[j];
        }
    }

//Сбор результирующего вектора
    if(rank == 0){
        for(int i = 0; i < N/size; i++)
            R[i] = localresult[i];
        for(int i = 1; i < size; i++){
            MPI_Recv(R+((N/size)*(i)), (N)/size + ((i < mod)? 1 : 0), MPI_DOUBLE, i, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }
    }
    else{
       MPI_Send(localresult, N/size + ((rank < mod)? 1 : 0), MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
    }
}

int Minimal_Nevazki(int N, double *A, const double *b, double *X, double eps)
//N - размерность, A - матрица, b - вектор свободных членов, X - вектор результат, eps - точность
{
    int count = 0;//Число итераций
    double *R = new double[N];
    double *Y = new double[N];
    double *Xn = new double[N];//приближение

    double crit_module;
    double chisl_Tau;
    double del_Tau;

    for(int i = 0; i < N; i++){
        Xn[i] = 0;//Начальное приближение делаем равным 0
    }

    do{
        Matrix_by_vector(N, A, Xn, R);
        for(int i = 0; i < N; i++){
            Y[i] = R[i] - b[i];
        }

        Matrix_by_vector(N, A, Y, R);

        chisl_Tau = 0.0;
        del_Tau = 0.0;
        for(int i = 0; i < N; i++)
        {
            chisl_Tau += R[i]*Y[i];
            del_Tau += R[i]*R[i];
        }
        chisl_Tau = chisl_Tau/del_Tau;
        for(int i = 0; i < N; i++){
            X[i] = Xn[i] - chisl_Tau*Y[i];
        }


        Matrix_by_vector(N, A, X, R);
        double crit_1 = 0.0;
        double crit_2 = 0.0;
        for(int i = 0; i < N; i++){
            crit_1 += pow(R[i] - b[i], 2);
            crit_2 += pow(b[i], 2);
        }
        crit_1 = sqrt(crit_1);
        crit_2 = sqrt(crit_2);
        crit_module = crit_1/crit_2;

        //Обновляем приближение
        for(int i = 0; i < N; i++){
            Xn[i] = X[i];
        }

        count++;
    }
    while (crit_module >= eps);

    delete[](R);
    delete[](Y);
    delete[](Xn);
    return count;
}

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

    int rank, size;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    int N = atoi(argv[1]);
    double A[N][N];
    double u[N];

//Нулевой процесс заполняет матрицу A и начальный вектор u
if(rank == 0)
{
    //Заполнение матрицы А
    for(int i = 0; i < N; i++){
        for(int j = 0; j < N; j++){
            if(i == j){
                A[i][j] = 2.0;
            } else{
                A[i][j] = 1.0;
            }
        }
    }

    for(int i = 0; i < N; i++){
        u[i] = sin((2*M_1_PI*i)/N);
    }
}

    double b[N];
    Matrix_by_vector(N, (double *)A, u, b);

    double X[N];//Вектор решений
    double epsilon = pow(10, -5);//Точность

    //std::cout << "Curr epsilon:" << epsilon << std::endl;

    double timer = MPI_Wtime();
    int count_steps = Minimal_Nevazki(N,(double *)A, b, X, epsilon);
    timer = MPI_Wtime() - timer;

    if(rank == 0){
      std::cout << "Count steps:" << count_steps << std::endl;
      std::cout << "Time:" << timer << std::endl;
      for(int i = 0; i < N; i++){
         std::cout << "X[" << i << "]:" << X[i] << "  U[" << i << "]:" << u[i] << std::endl;
      }
    }

    MPI_Finalize();
    return 0;
}