#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;
}
