! Distributing a 4x4 matrix blockwise,
! and then using user-defined communicators and
! reduction to sum over first index of old matrix

program comms
  use mpi_f08
  implicit none

  integer :: i,rank,nproc
  integer :: clr,key,rank_c1,nproc_c1
  real(kind(1d0)) :: full(4,4), corner(2,2),s,s2
  type(mpi_comm) :: comm1,comm2
  type(mpi_status) :: stat

  call mpi_init()
  call mpi_comm_size(mpi_comm_world, nproc)
  call mpi_comm_rank(mpi_comm_world, rank)

  if ((rank==0).and.(nproc<4)) then
    write(*,*)'Too few ranks: ',nproc
    call mpi_abort(mpi_comm_world,1)
  endif

  call random_number(full)

  if (rank.eq.0) then
    corner=full(1:2,1:2)
    call mpi_send(full(1:2,3:4),4,MPI_DOUBLE_PRECISION,1,1,mpi_comm_world)
    call mpi_send(full(3:4,1:2),4,MPI_DOUBLE_PRECISION,2,1,mpi_comm_world)
    call mpi_send(full(3:4,3:4),4,MPI_DOUBLE_PRECISION,3,1,mpi_comm_world)
  else
    call mpi_recv(corner,4,MPI_DOUBLE_PRECISION,0,1,mpi_comm_world,stat)
  endif

  if (rank==0) then
    write(*,*)'On root sums:'
    do i=1,4
      write(*,'(i5,f8.4)')i,sum(full(:,i))
    enddo
  endif

  clr=mod(rank,2)
  key=1
  if (rank>=2) key=2
!  write(*,*)'Rank ',rank,' has colour ',clr,' and key ',key
  call mpi_comm_split(mpi_comm_world,clr,key,comm1)
  call mpi_comm_rank(comm1, rank_c1)
  call mpi_comm_size(comm1, nproc_c1)

!  write(*,*)'Old rank',rank,'is',rank_c1,'of',nproc_c1,'in comm1'

  if (rank==0) write(*,*)'Distributed sums'
  do i=1,2
     s=sum(corner(:,i))
     call mpi_reduce(s,s2,1,MPI_DOUBLE_PRECISION,MPI_SUM,0,comm1)
     if (rank_c1==0) write(*,'(2i5,f8.4)')rank,i,s2
   enddo

  call mpi_finalize()

end program
