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

int mat_dim = 16;
int blk_dim = 4;
MPI_Win win;

void *thread_work(void *arg)
{
    int i;
    double *local_b;
    MPI_Alloc_mem(sizeof(double) * blk_dim * blk_dim, MPI_INFO_NULL, &local_b);

    for (i = 0; i < 1000; i++) {
        MPI_Get(local_b, blk_dim * blk_dim, MPI_DOUBLE,
                0, 0, blk_dim * blk_dim, MPI_DOUBLE, win);

        MPI_Win_flush_all(win);
    }

    MPI_Free_mem(local_b);

    pthread_exit(NULL);
    return NULL;
}

int main(int argc, char **argv)
{
	int rank, nprocs, provided;
    int i, num_threads;
    double *win_mem;
    pthread_t *threads;
    void *status;

    /* initialize MPI environment */
	MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided);
	if (provided < MPI_THREAD_MULTIPLE) {
		MPI_Abort(MPI_COMM_WORLD, 1);
		exit(-1);
	}
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

    if (argc < 2) {
        printf("usage: %s <t>\n", argv[0]);
        MPI_Finalize();
        exit(-1);
    }
    num_threads = atoi(argv[1]);
    threads = (pthread_t *)malloc(num_threads * sizeof(pthread_t));

    if (rank == 0) {
        printf("num_threads: %d\n", num_threads);

        MPI_Win_allocate(mat_dim * mat_dim * sizeof(double), sizeof(double),
                         MPI_INFO_NULL, MPI_COMM_WORLD, &win_mem, &win);
    } else {
        MPI_Win_allocate(0, sizeof(double), MPI_INFO_NULL, MPI_COMM_WORLD,
                         &win_mem, &win);
    }

    MPI_Barrier(MPI_COMM_WORLD);

    MPI_Win_lock_all(0, win);

    for (i = 0; i < num_threads; i++) {
        pthread_create(&threads[i], NULL, thread_work, NULL);
    }
    for (i = 0; i < num_threads; i++) {
        pthread_join(threads[i], &status);
    }

    MPI_Win_unlock_all(win);

    MPI_Barrier(MPI_COMM_WORLD);

    if (rank == 0) {
        printf("Done.\n");
    }

    free(threads);

    MPI_Win_free(&win);
	MPI_Finalize();

	return EXIT_SUCCESS;
}