program test_mpiio
! Test of writing MPI subarray datatype using MPI-IO. Each rank holds a plane of 
! a larger (m x n x n)-sized integer and double arrays, where m is the total number of ranks.

   use mpi
   implicit none

   integer, parameter :: n = 100 ! Number of columns
   integer :: m, i, myrank, fullsize(3), subsize(3), subind(3), ierr, subarray, subxarray, fh, a(n,n,3)
   double precision :: x(n,n,3),y(n,n,6)
   integer(kind=mpi_offset_kind) :: disp = 0

   call mpi_init(ierr)
   call mpi_comm_rank( mpi_comm_world, myrank, ierr )
   call mpi_comm_size( mpi_comm_world, m, ierr )      ! m is the number of ranks

   a = myrank ! Set my row to value of my rank
   x = myrank + 0.1
   y = -1.0

   fullsize = [n, m, n] ! Size of full array in file

   if (m >= 3) then
     if (myrank == 0 .or. myrank == m-1) then
       subsize  = [n, 1, n]  ! Size of one plane held by my rank at the beginning and end
       subind   = [0,myrank,0] ! Starting indices of my rank's plane
     else
       subsize = [n, 3, n]  ! Overlapping planes in the middle of the data
       subind = [0, myrank - 1, 0]
     end if
   else
     write(*,*) "Run on at last 3 processors please."
     stop
   end if

   ! Create subarry of row held by my rank from full array
   call mpi_type_create_subarray(3, fullsize, subsize, subind, mpi_order_fortran, mpi_integer, subarray, ierr)
   call mpi_type_commit(subarray, ierr )
   call mpi_type_create_subarray(3, fullsize, subsize, subind, mpi_order_fortran, mpi_double, subxarray, ierr)
   call mpi_type_commit(subxarray, ierr )

   ! Open file, set view to subarray, write contents, and close file
   call mpi_file_open(mpi_comm_world, 'array.dat', mpi_mode_wronly+mpi_mode_create, mpi_info_null, fh, ierr)
   call mpi_file_set_view(fh, disp, mpi_integer, subarray, 'native', mpi_info_null, ierr)
   if (myrank == 0 .or. myrank == m-1) then
     call mpi_file_write_all(fh, a, n*n, mpi_integer, mpi_status_ignore, ierr)
   else
     call mpi_file_write_all(fh, a, n*n*3, mpi_integer, mpi_status_ignore, ierr)
   end if

   disp = n*m*n*4
   call mpi_file_set_view(fh, disp, mpi_double, subxarray, 'native', mpi_info_null, ierr)
   if (myrank == 0 .or. myrank == m-1) then
     call mpi_file_write_all(fh, x, n*n, mpi_double, mpi_status_ignore, ierr)
     call mpi_file_write_all(fh, x, n*n, mpi_double, mpi_status_ignore, ierr)
   else
     call mpi_file_write_all(fh, x, n*n*3, mpi_double, mpi_status_ignore, ierr)
     call mpi_file_write_all(fh, x, n*n*3, mpi_double, mpi_status_ignore, ierr)
   endif
   call mpi_file_close(fh, ierr)

   ! Read the file
   disp = 0
   a = -1
   x = -1.0

   call mpi_file_open(mpi_comm_world, 'array.dat', mpi_mode_rdonly, mpi_info_null, fh, ierr)
   call mpi_file_set_view(fh, disp, mpi_integer, subarray, 'native', mpi_info_null, ierr)
   if (myrank == 0 .or. myrank == m-1) then
     call mpi_file_read_all(fh, a, n*n, mpi_integer, mpi_status_ignore, ierr)
   else
     call mpi_file_read_all(fh, a, n*n*3, mpi_integer, mpi_status_ignore, ierr)
   end if
 
   disp = n*m*n*4
   call mpi_file_set_view(fh, disp, mpi_double, subxarray, 'native', mpi_info_null, ierr)
   if (myrank == 0 .or. myrank == m-1) then
     call mpi_file_read_all(fh, y, n*n*2, mpi_double, mpi_status_ignore, ierr)
   else
     call mpi_file_read_all(fh, y, n*n*3*2, mpi_double, mpi_status_ignore, ierr)
   end if
   call mpi_file_close(fh, ierr)

   if (m < 6) then
     if (myrank == 0 .or. myrank == m-1) then
       write(*,'(*(I4))') myrank,a(:,:,1)
       write(*,'(I4,2x,*(F5.1))') myrank,y(:,:,1:2)
     else
       write(*,'(*(I4))') myrank,a
       write(*,'(I4,2x,*(F5.1))') myrank,y
     end if
   end if

   call mpi_finalize(ierr)
   stop

end program test_mpiio
