#include <mpi.h>
#include <iostream>
#include <vector>

#include <pnetcdf.h>

inline void check_ncmpi_status(MPI_Comm comm,
                               int stat, int ncid, const char* msg=NULL)
{
  if ( stat ) {
    if ( msg ) 
      std::cout << msg << ": ";
    std::cout << "Parallel NetCDF error! stat=" << stat << std::endl;
    ncmpi_abort(ncid);
    MPI_Abort(comm, -1);
  }
}


int main(int argc, char* argv[])
{
  MPI_Init(&argc, &argv);
  MPI_Comm comm = MPI_COMM_WORLD;
  int size;
  MPI_Comm_size(comm, &size);
  int rank;
  MPI_Comm_rank(comm, &rank);
    
  {
    int cmode = NC_64BIT_OFFSET;
    const char* file = "test.ncdf";
    int stat;
    int ncid;
    MPI_Info info;
    MPI_Info_create(&info);
    MPI_Info_set(info, "romio_ds_read", "disable");

    stat = ncmpi_create(comm, file, cmode, info, &ncid);
    check_ncmpi_status(comm, stat, ncid, "create");
    
    int  length = 500000000;    
    //int  length = 500;
    MPI_Offset count = length / size;
    //make sure length is correct
    length = count * size;
    MPI_Offset start = rank * count;

    //dimension
    const char* dim_name = "length";
    int dimid;
    stat = ncmpi_def_dim(ncid, dim_name, length, &dimid);
    check_ncmpi_status(comm, stat, ncid, "define dim");

    //variable
    int vid;
    stat = ncmpi_def_var(ncid, "eigenvector", NC_DOUBLE, 1, &dimid, &vid);
    check_ncmpi_status(comm, stat, ncid, "define var");

    //enddef
    stat = ncmpi_enddef(ncid);
    check_ncmpi_status(comm, stat, ncid, "end def");  
    
    std::vector<double> v(count, 1.0);

    if ( rank == 0 ) std::cout << " Writing data ...";
    //write data
    stat = ncmpi_put_vara_double_all(ncid, vid,
                                     &start, &count, &v[0]);
    check_ncmpi_status(comm, stat, ncid, "write");
    stat = ncmpi_close(ncid);
    check_ncmpi_status(comm, stat, ncid, "close");
    if ( rank == 0 ) std::cout << "  ... Done" << std::endl;

    MPI_Info_delete(info, "romio_ds_read");
    MPI_Info_free(&info);
  }

  MPI_Finalize();
  return 0;
}