#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>
#include "tools.c"

#define NDIM      2
#define NX_GLOB   16
#define NY_GLOB   8

#define VERSION   2

int main(int argc, char ** argv)
{
  int i, j, rank, size;
  int nx, ny;
  int nprocs[NDIM];
  int gsizes[NDIM];
  int lsizes[NDIM];
  int periods[NDIM] = {0,0};
  int coords[NDIM];
  int start[NDIM];
  double **A;
  MPI_Datatype subarr_type;
  MPI_Comm MPI_COMM_CART;
  char fname[] = "arr2D.bin";

/* --------------------------------------------------------
   0. Initialize the MPI execution environment
   -------------------------------------------------------- */

  MPI_Init (&argc, &argv);
  MPI_Comm_rank (MPI_COMM_WORLD, &rank);
  MPI_Comm_size (MPI_COMM_WORLD, &size);

/* --------------------------------------------------------
   1. Create a 2D domain decomposition
   -------------------------------------------------------- */
  
  /* -- 1a. Attempt to find a maximally cubic decomposition -- */

  nprocs[0] = (int)sqrt(size);
  nprocs[1] = size/nprocs[0];
  if (nprocs[0]*nprocs[1] != size){
    if (rank == 0) printf ("! Cannot decompose\n");
    MPI_Finalize();
    return 1;
  }

  /* -- 1b. Create communicator -- */

  MPI_Cart_create(MPI_COMM_WORLD, NDIM, nprocs, periods, 0, &MPI_COMM_CART);
  MPI_Cart_get(MPI_COMM_CART, NDIM, nprocs, periods, coords);

  gsizes[0] = NX_GLOB;
  gsizes[1] = NY_GLOB;

  lsizes[0] = nx = NX_GLOB/nprocs[0];
  lsizes[1] = ny = NY_GLOB/nprocs[1];

  if (rank == 0){
    printf ("Domain decomposed in %d X %d procs\n",nprocs[0],nprocs[1]);
    printf ("Local grid size = %d X %d\n",lsizes[0], lsizes[1]);
  }

/* --------------------------------------------------------
   2. Allocate memory and fill 2D array on local domain
   -------------------------------------------------------- */

  A = Allocate_2DdblArray(ny,nx);
  for (j = 0; j < ny; j++) {
  for (i = 0; i < nx; i++) {
    A[j][i] = rank;
  }}

/* --------------------------------------------------------
   3. Create new datatypes
   -------------------------------------------------------- */

  start[0] = coords[0]*lsizes[0];
  start[1] = coords[1]*lsizes[1];

  MPI_Type_create_subarray (NDIM, gsizes, lsizes, start, MPI_ORDER_FORTRAN,
                            MPI_DOUBLE, &subarr_type);
  MPI_Type_commit (&subarr_type);

/* --------------------------------------------------------
   4. Open file for writing
   -------------------------------------------------------- */

  MPI_File_delete(fname, MPI_INFO_NULL);

  MPI_File fh;
  MPI_Status status;

  MPI_File_open(MPI_COMM_CART, fname, 
                MPI_MODE_CREATE | MPI_MODE_WRONLY, MPI_INFO_NULL, &fh);
  MPI_File_set_view(fh, 0, MPI_DOUBLE, subarr_type, "native", MPI_INFO_NULL);
  MPI_File_write_all(fh, A[0], nx*ny, MPI_DOUBLE, MPI_STATUS_IGNORE);
  MPI_File_close(&fh);

  MPI_Type_free(&subarr_type);
  MPI_Finalize();
  return 0;
}