#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[30] = "Collective write example";
  int dims[2];

  MPI_Init(&argc, &argv);

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

  // Make sure that dim1 is evenly divisble by 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]);


  // creating a new file  (is in define mode by default)
  ret = ncmpi_create(MPI_COMM_WORLD, 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", 30, buf);
  if (ret != NC_NOERR) err(ret, __LINE__);

  // switch to data mode
  ret = ncmpi_enddef(ncid);
  if (ret != NC_NOERR) err(ret, __LINE__);


  MPI_Offset start[2], count[2];
  int *data = (int*)malloc(sizeof(int)*sd1*sd2/nprocs);
  int i;
  for (i = 0; i < sd1*sd2 / nprocs; i++)
   data[i] = rank+1;

  start[0] = (sd1/nprocs)  * rank;
  start[1] =  0;

  count[0] = sd1/nprocs;
  count[1] = sd2;

  double s, e; 
  s = MPI_Wtime(); 

  // store variables in collective (default)
  ret = ncmpi_put_vara_int_all(ncid, varid1, start, count, data);
  if (ret != NC_NOERR) err(ret, __LINE__);
  
  // close the file 
  ncmpi_close(ncid);

  e = MPI_Wtime();

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

  free(data);

  MPI_Finalize();
  return 0;
}