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

static void err(int status, int lineno)
{
    fprintf(stderr, "Error at line %d: %s\n", lineno, ncmpi_strerror(status));
    MPI_Abort(MPI_COMM_WORLD, 1);
}

int main(int argc, char **argv) {
  int ret, ncid=0, nprocs, rank, dimid1, dimid2, varid1=0, varid2=0, ndims=2;
  int sd1, sd2;
  char buf[50] = "Single process write example";
  int dims[2];

  MPI_Init(&argc, &argv);

  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

  if (argc != 4) {
      if (rank == 0) printf("Usage: %s <filename> <dim1 of a 2D array> <dim2>\n", argv[0]);
      MPI_Finalize();
      exit(-1);
  }

  sd1 = atoi(argv[2]);
  sd2 = atoi(argv[3]);

  int *data = NULL;
  int *localdata = (int *) malloc(sizeof(int)*sd1*sd2/nprocs);
  if (rank == 0) {
      ret = ncmpi_create(MPI_COMM_SELF, argv[1],NC_CLOBBER|NC_64BIT_OFFSET, MPI_INFO_NULL, &ncid);
      if (ret != NC_NOERR) err(ret, __LINE__);


      // Defining the file contents - dimensions, variables, attributes
      ret = ncmpi_def_dim(ncid, "d1", sd1, &dimid1);
      if (ret != NC_NOERR) err(ret, __LINE__);
      ret = ncmpi_def_dim(ncid, "d2", sd2, &dimid2);
      dims[0] = dimid1; dims[1] = dimid2;
      if (ret != NC_NOERR) err(ret, __LINE__);
      ret = ncmpi_def_var(ncid, "v", NC_INT, ndims, dims, &varid1);
      if (ret != NC_NOERR) err(ret, __LINE__);
    
      // add attributes
      ret = ncmpi_put_att_text(ncid, NC_GLOBAL, "string", 50, buf);
      if (ret != NC_NOERR) err(ret, __LINE__);
    
      // switch to data mode
      ret = ncmpi_enddef(ncid);
      if (ret != NC_NOERR) err(ret, __LINE__);

      data = (int*)malloc(sizeof(int)*sd1*sd2);
  }
  int i;
  for (i = 0; i < sd1*sd2 / nprocs; i++)
   localdata[i] = rank+1;;
  
  double s,e;
  s = MPI_Wtime();
  MPI_Gather(localdata, sd1*sd2 / nprocs, MPI_INT, data, sd1*sd2 / nprocs, MPI_INT, 0, MPI_COMM_WORLD);

  if (rank == 0) {
      MPI_Offset start[2], count[2];
      start[0] = start[1] = 0;
      count[0] = sd1; count[1] = sd2;
      
      ret = ncmpi_put_vara_int_all(ncid, varid1, start, count, data);
      if (ret != NC_NOERR) err(ret, __LINE__);

      ret = ncmpi_close(ncid);
      if (ret != NC_NOERR) err(ret, __LINE__);

      free(data);
  }
  e = MPI_Wtime();

  if (rank == 0)
   printf("Output time = %f\n", e-s);

  free(localdata);

  MPI_Finalize();
  return 0;
}