#include <stdio.h>
#include <math.h>
#include <mpi.h>
#include <stdlib.h>
#include <string.h>

#define N 1000
#define START (-100)
#define STOP (100)
#define ITERATIONS_COUNT 10000

void calculateRightPart(int rank, const int *b_send_counts, const int *b_displs, const double *b, double *right_part) {
    double *start_b = (double *)malloc(b_send_counts[rank] * sizeof(double));
    MPI_Scatterv(b, b_send_counts, b_displs, MPI_DOUBLE, start_b, b_send_counts[rank], MPI_DOUBLE, 0, MPI_COMM_WORLD);
    size_t local_sum = 0;
    size_t global_sum = 0;
    for (int i = 0; i < b_send_counts[rank]; i++) {
        local_sum += start_b[i];
    }
    MPI_Reduce(&local_sum, &global_sum, 1, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);

    double epsilon = powl(10, -5);

    if(rank == 0){
        (*right_part) = global_sum * epsilon * epsilon;
    }
    MPI_Bcast(right_part, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    free(start_b);
}


double * generate_double_arr(size_t n){
    return (double *) malloc(n * sizeof(double));
}

int * generate_int_arr(size_t n){
    return (int *) malloc(n * sizeof(int));
}

double* read_matrix_from_file(const char* filename) {
    FILE* fp = fopen(filename, "r");
    if (fp == NULL) {
        printf("Error opening file %s\n", filename);
        return NULL;
    }

    size_t n;
    fscanf(fp, "%zd", &n);
    char c = fgetc(fp);

    // Allocate memory for the matrix
    double* arr = generate_double_arr(n*n);

    // Read the matrix elements from the file
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            double res;
            fscanf(fp, "%lf", &res);
            arr[i* n + j] = res;
            fgetc(fp);
        }
    }

    fclose(fp);
    return arr;
}

double* read_vector_from_file(const char* filename) {
    FILE* fp = fopen(filename, "r");
    if (fp == NULL) {
        printf("Error opening file %s\n", filename);
        return NULL;
    }

    size_t n;
    fscanf(fp, "%zd", &n);
    fgetc(fp);

    // Allocate memory for the matrix
    double* arr = generate_double_arr(n);

    // Read the matrix elements from the file
    for (int i = 0; i < n; i++) {
        double res;
        fscanf(fp, "%lf", &res);
        arr[i] = res;
        fgetc(fp);
    }

    fclose(fp);
    return arr;
}



void fill_vector_with_value(double *res, double value) {
    for (size_t i = 0; i < N; i++) {
        res[i] = value;
    }
}


double multiple_row_and_column_vector(double *row, double *column, size_t size){
    double res = 0;
    for(size_t i = 0; i < size; i++){
        res += row[i] * column[i];
    }
    return res;
}

void print_vec(double *vec, size_t size){
    for(size_t i = 0; i < size; i++){
        printf("%lf\n", vec[i]);
    }
    printf("\n\n");
}

void subtract_vectors(const double *vec1, const double *vec2, double *res, int size){
    for(size_t i = 0; i < (size_t)size; i++){
        res[i] = vec1[i] - vec2[i];
    }
}

void multiple_matrix_and_vector(const double *matrix, const double *vec, double *res){
    memset(res, 0, sizeof(double) * N);
    for(size_t i = 0; i < N; i++) {
        for (size_t j = 0; j < N; j++) {
            res[j] += matrix[j*N + i] * vec[i];
        }
    }

}

void calculate_vector_y(const double *matrix_a, double* x, double *b, double *res) {
    multiple_matrix_and_vector(matrix_a, x, res);
    subtract_vectors(res, b, res, N);
}

int main(int argc, char ** argv){
    int rank, size, n, rest;
    MPI_Init( & argc, & argv);
    MPI_Comm_size(MPI_COMM_WORLD, & size);
    MPI_Comm_rank(MPI_COMM_WORLD, & rank);

    n = N / size;
    rest = N % size;

    int * matrix_send_counts = generate_int_arr(size);
    int *  matrix_displs = generate_int_arr(size);

    int *  b_send_counts = generate_int_arr(size);
    int * b_displs = generate_int_arr(size);

    n = N / size;
    rest = N % size;

    int offset = 0;
    for (int i = 0; i < size; i++) {
        matrix_send_counts[i] = N / size * N;
        if (i < rest) {
            matrix_send_counts[i] += N;
        }
        matrix_displs[i] = offset;
        offset += matrix_send_counts[i];
    }



    offset = 0;
    for (int i = 0; i < size; i++) {
        b_send_counts[i] = N / size;
        if (i < rest) {
            b_send_counts[i] += 1;
        }
        b_displs[i] = offset;
        offset += b_send_counts[i];
    }

    double *a, *b, *x, *y, *tau_y;
    x = generate_double_arr(N);
    if(rank == 0){
        a = read_matrix_from_file("matrix.txt");
        b = read_vector_from_file("vector.txt");
        fill_vector_with_value(x, 0);
        y = generate_double_arr(N);
        tau_y = generate_double_arr(N);
    }
    double right_part;
    calculateRightPart(rank, b_send_counts, b_displs, b, &right_part);

    if(rank == 0){
        printf("%f", right_part);
    }

    double *matrix_rec_row = generate_double_arr(matrix_send_counts[rank]);
    double *vec_rec1 = generate_double_arr(b_send_counts[rank]);

    MPI_Scatterv(a, matrix_send_counts, matrix_displs, MPI_DOUBLE, matrix_rec_row, matrix_send_counts[rank], MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(x, N, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Scatterv(b, b_send_counts, b_displs, MPI_DOUBLE, vec_rec1, b_send_counts[rank], MPI_DOUBLE, 0, MPI_COMM_WORLD);

    int num_handle_rows = matrix_send_counts[rank] / N;

    //cout << num_handle_rows << "\n";


    for(size_t i = 0; i < num_handle_rows; i++){
        vec_rec1[i] = multiple_row_and_column_vector(&matrix_rec_row[num_handle_rows*i], x, b_send_counts[rank]) - vec_rec1[i*num_handle_rows];
        //cout << b[i*N] << "\n";
    }

    MPI_Gatherv(vec_rec1, num_handle_rows, MPI_DOUBLE, y, b_send_counts,  b_displs, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    if(rank == 0){
        print_vec(y, 10);
        double *kek = generate_double_arr(N);
        calculate_vector_y(a, x, b, kek);
        print_vec(kek, 10);
        free(kek);
    }

    free(matrix_send_counts);
    free(b_send_counts);
    free(matrix_displs);
    free(b_displs);
    free(matrix_rec_row);
    free(vec_rec1);
    if(rank == 0){
        free(a);
        free(b);
        free(y);
        free(tau_y);
    }
    MPI_Finalize();
}