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

void matrix_vector_multiply(double *A, double *b, double *c, int start_row, int end_row, int M) {
    for (int i = start_row; i < end_row; i++) {
        c[i] = 0;
        for (int j = 0; j < M; j++) {
            c[i] += A[i * M + j] * b[j];
        }
    }
}

int parse_integer(const char *str) {
    int value = 0;
    while (*str) {
        if (*str < '0' || *str > '9') {
            return -1; // Invalid input
        }
        value = value * 10 + (*str - '0');
        str++;
    }
    return value;
}

int main(int argc, char **argv) {
    int rank, size, N, M;

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

    // Check command-line arguments for matrix dimensions
    if (argc != 3) {
        if (rank == 0) {
            printf("Usage: %s <num_rows> <num_columns>\n", argv[0]);
        }
        MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
    }


    // Parse matrix dimensions
    N = parse_integer(argv[1]);
    
    M = parse_integer(argv[2]);

    if (N <= 0 || M <= 0) {
        if (rank == 0) {
            printf("Invalid matrix dimensions: %s x %s\n", argv[1], argv[2]);
        }
        MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
    }

    double *A = NULL;
    double *b = NULL;
    double *c = (double *)malloc(N * sizeof(double));

    // Master process initializes the matrix and vector
    if (rank == 0) {
        A = (double *)malloc(N * M * sizeof(double));
        b = (double *)malloc(M * sizeof(double));

        // Initialize matrix A and vector b with simple values
        for (int i = 0; i < N; i++)
            for (int j = 0; j < M; j++)
                A[i * M + j] = 1.0;  // Example initialization
        for (int j = 0; j < M; j++)
            b[j] = 1.0;  // Example initialization
    }

    // Broadcast vector b to all processes
    MPI_Bcast(b, M, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    // Each process computes its assigned rows
    int rows_per_process = N / size;
    int start_row = rank * rows_per_process;
    int end_row = (rank == size - 1) ? N : start_row + rows_per_process;

    double *local_c = (double *)malloc((end_row - start_row) * sizeof(double));
    matrix_vector_multiply(A, b, local_c, start_row, end_row, M);

    // Gather results in the master process
    MPI_Gather(local_c, end_row - start_row, MPI_DOUBLE, c, end_row - start_row, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    // Master process prints the result
    if (rank == 0) {
        for (int i = 0; i < N; i++) {
            printf("c[%d] = %f\n", i, c[i]);
        }
        free(A);
        free(b);
    }

    free(local_c);
    free(c);
    MPI_Finalize();
    return 0;

    
}