#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
using namespace std;

int multiplyRowByVector(int *row, int *vector, int size, int rank) {
  int sum = 0;
  cout << "Rank: " << rank << " row: " << row[0] << row[1] << row[2] << endl;
  cout << "Rank: " << rank << " vec: " << vector[0] << vector[1] << vector[2] << endl;
  for (int i = 0; i < size; i++) {
    sum = sum + row[i] * vector[i];
  }
  return sum;
}

int receiveAndMultiply(int rank, int size) {
  int *row = new int[size];
  int *vector = new int[size];
  int result = 0;

  MPI_Recv(row, size, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
  MPI_Recv(vector, size, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
  result = multiplyRowByVector(row, vector, size, rank);
  MPI_Send(&result, 1, MPI_INT, 0, 0, MPI_COMM_WORLD);
}

int** createMatrix(int size) {
  int** matrix = new int*[size];
  for (int i = 0; i < size; ++i) {
    matrix[i] = new int[size];
    for (int j = 0; j < size; j++) {
      matrix[i][j] = rand() % 9 + 1;
    }
  }
  return matrix;
}

int* createVector(int size) {
  int* vector = new int[size];
  for (int i = 0; i < size; i++) {
    vector[i] = rand() % 9 + 1;
  }
  return vector;
}

void drawMatrix(int** matrix, int size) {
  for (int i = 0; i < size; i++) {
    for (int j = 0; j < size; j++) {
      cout << matrix[i][j];
    }
  cout << endl;
  }
}

int main(int argc, char** argv) {
  int rank = 0;
  int size = 0;
  int result_from_proc = 0;
  int* results = new int[size];
  MPI_Init(NULL, NULL);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);

  if (rank == 0) {
    int** matrix = createMatrix(size);
    int* vector = createVector(size);
    drawMatrix(matrix, size);
    results[0] =  multiplyRowByVector(matrix[0], vector, size, rank);
    for (int i = 1; i < size; i++) {
      MPI_Send(matrix[i], size, MPI_INT, i, 0, MPI_COMM_WORLD);
      MPI_Send(vector, size, MPI_INT, i, 0, MPI_COMM_WORLD);
      MPI_Recv(&result_from_proc, 1, MPI_INT, i, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      results[i] = result_from_proc;
    }
  } else {
    receiveAndMultiply(rank, size);
  }

  if (rank == 0) {
  cout << "Results: " << endl;
    for (int i = 0; i < size; i++) {
       cout << "Rank: " << i << " " << results[i] << endl;
    }
  }

  MPI_Finalize();
}