PROGRAM array
!This code domain splits a 2-D grid and solves the Heat equation using 
!Explicit Euler.
IMPLICIT NONE

INCLUDE "mpif.h"
INCLUDE "declarationsAppendix.inc"
!set up the MPI stuff with cart grid etc
CALL MPISETUP(rank,nproc_x,nproc_y,comm,coords,nx_global,ny_global,nx,ny,   &
               left,right,up,down)

!Set the initial conditions and communicate amongst the slaves
IF (rank == 0) THEN
  ALLOCATE(x_glob(nx_global),y_glob(ny_global))
  ALLOCATE(IC_glob(nx_global,ny_global))
  dx=1000./nx_global ;  dy=1000./ny_global
  !define the spatial mesh and initial condition
  CALL GETGLOBIC(nx_global,ny_global,dx,dy,x_glob,y_glob,IC_glob)

!write out your solution
open(unit=90,file='TEST.dat')
do ix=1,nx_global
do iy=1,ny_global
write(90,'(3f16.4)')x_glob(ix),y_glob(iy),IC_glob(ix,iy)
enddo
write(90,'(a)')""
enddo
close(90)
ENDIF
!Set up the local spatial mesh and initial condition arrays
ALLOCATE(x(0:nx+1),y(0:ny+1),uold(0:nx+1,0:ny+1),unew(0:nx+1,0:ny+1))

!Pass read in data out to slave nodes etc.
CALL PASSINFO(nx,ny,nx_global,ny_global,nproc_x,nproc_y,rank,nextrank,comm, &
                    x,x_glob,y,y_glob,dx,dy,uold,IC_glob)

!Clean up
if(rank == 0)then 
!IF (ALLOCATED(x_glob)) DEALLOCATE(x_glob)
!IF (ALLOCATED(y_glob)) DEALLOCATE(y_glob)
IF (ALLOCATED(IC_glob)) DEALLOCATE(IC_glob)
endif

!allocate the sizes of the strips that will be communicated between processes
ALLOCATE(works1(1:ny),works2(1:ny),workr1(1:ny),workr2(1:ny))

!begin time loop
do 
  if(rank==0) then
  print*,t
  endif
  
  !communicate strips to the neighbouring processors in horizontal direction
  CALL MPI_ISEND(uold(1,ny),nx,MPI_REAL,down,tag,MPI_COMM_WORLD,jsend1,ierr)
  CALL MPI_ISEND(uold(1,1),nx,MPI_REAL,up,tag,MPI_COMM_WORLD,jsend2,ierr)
  CALL MPI_IRECV(uold(1,0),nx,MPI_REAL,up,tag,MPI_COMM_WORLD,jrecv1,ierr)
  CALL MPI_IRECV(uold(1,ny+1),nx,MPI_REAL,down,tag,MPI_COMM_WORLD,jrecv2,ierr)
  CALL MPI_WAIT(jsend1,stat,ierr)
  CALL MPI_WAIT(jsend2,stat,ierr)
  CALL MPI_WAIT(jrecv1,stat,ierr)
  CALL MPI_WAIT(jrecv2,stat,ierr)
  
  !communicate strips to the neighbouring processors in vertical direction
  if(left.ne.MPI_PROC_NULL)then
    do iy=1,ny
      works1(iy)=uold(1,iy)
    enddo
  endif
  if(right.ne.MPI_PROC_NULL)then
    do iy=1,ny
      works2(iy)=uold(nx,iy)
    enddo
  endif
  CALL MPI_ISEND(works1(1),ny,MPI_REAL,left,tag,comm,jsend1,ierr)
  CALL MPI_ISEND(works2(1),ny,MPI_REAL,right,tag,comm,jsend2,ierr)
  CALL MPI_IRECV(workr1(1),ny,MPI_REAL,right,tag,comm,jrecv1,ierr)
  CALL MPI_IRECV(workr2(1),ny,MPI_REAL,left,tag,comm,jrecv2,ierr)
  CALL MPI_WAIT(jsend1,stat,ierr)
  CALL MPI_WAIT(jsend2,stat,ierr)
  CALL MPI_WAIT(jrecv1,stat,ierr)
  CALL MPI_WAIT(jrecv2,stat,ierr)
  
  if(left.ne.MPI_PROC_NULL)then
    do iy=1,ny
      uold(0,iy)=workr2(iy)
    enddo
  endif
  if(right.ne.MPI_PROC_NULL)then
    do iy=1,ny
      uold(nx+1,iy)=workr1(iy)
    enddo
  endif
  
  
  !update u array using Explicit Euler
  do ix=1,nx
    do iy=1,ny
  !boundary conditions are simply N=0 at the boundaries
      if(ix==1.and.left==MPI_PROC_NULL)then
        unew(ix,iy)=0.0
      else if(ix==nx.and.right==MPI_PROC_NULL)then
        unew(ix,iy)=0.0
      else if(iy==1.and.up==MPI_PROC_NULL)then
        unew(ix,iy)=0.0
      else if(iy==ny.and.down==MPI_PROC_NULL)then
        unew(ix,iy)=0.0
      else         
        !Do the explicit Euler step
  unew(ix,iy)=uold(ix,iy)+dt*200*((uold(ix,iy+1)-2*uold(ix,iy)+uold(ix,iy-1))/(dy**2)+  &
      (uold(ix+1,iy)-2*uold(ix,iy)+uold(ix-1,iy))/(dx**2))  
      endif
    enddo
  enddo
  !update ready for next step
  t=t+dt
  uold=max(0.0,unew)
  
  !exit statement
  if(t.gt.finishtime)exit
!end time loop
enddo

!COLLECT SOLUTION FROM ALL CPUS
ALLOCATE(IC_glob(nx_global,ny_global))
CALL COLLECTINFO(nx,ny,nx_global,ny_global,nproc_x,nproc_y,rank,nextrank,comm,uold,IC_glob)

if(rank==0)then
!write out your solution
open(unit=90,file='TEST2.dat')
do ix=1,nx_global
do iy=1,ny_global
write(90,'(3f16.4)')x_glob(ix),y_glob(iy),IC_glob(ix,iy)
enddo
write(90,'(a)')""
enddo
close(90)

endif

CALL MPI_BARRIER(MPI_COMM_WORLD,ierr)
CALL MPI_FINALIZE(ierr)
END PROGRAM array

!*********************************************************************************************

subroutine MPISETUP(rank,nproc_x,nproc_y,comm,coords,nx_global,ny_global,nx,ny,   &
              left,right,up,down)
!This subrouting sets up the MPI grid of processes with nproc_x in the
!horizontal and nproc_y in the vertical direction. The subroutine should be
!appropriate for any 2-D grid splitting problem.
!nx_glob and ny_glob are the total dimensions of the 2-D computational grid
!nx and ny are the local size of the sub grid of cells being calculated by a
!particlar process
IMPLICIT NONE
INCLUDE "mpif.h"
INTEGER :: rank,nproc
INTEGER :: left,right,up,down
INTEGER :: ierr,ix,iy
INTEGER, DIMENSION(MPI_STATUS_SIZE) :: stat
INTEGER :: comm,ndims
INTEGER :: tag=100
INTEGER :: dims(2),coords(2)
INTEGER :: nx_global, ny_global
INTEGER :: nx, ny
INTEGER :: nproc_x,nproc_y
LOGICAL :: periods(2)

ndims=2
dims=(/0,0/)


!Start MPI
CALL MPI_INIT(ierr)
!Get total number of processes
CALL MPI_COMM_SIZE(MPI_COMM_WORLD,nproc,ierr)

!Create the 2D cartesian grid
!Returns dims which holds the x and y dim of the PE grid
CALL MPI_DIMS_CREATE(nproc,ndims,dims,ierr)
nproc_x=dims(1)
nproc_y=dims(2)
periods=.FALSE.
!creates a non periodic cartesian grid of CPU's 
CALL MPI_CART_CREATE(MPI_COMM_WORLD,2,dims,periods,.TRUE.,comm,ierr)
!Get the rank in the new communicator (does not have to be the same as in MPI_COMM_WORLD)
CALL MPI_COMM_RANK(comm,rank,ierr)
!Get nearest neighbours in 2D now
CALL MPI_CART_SHIFT(comm,0,1,left,right,ierr)
CALL MPI_CART_SHIFT(comm,1,1,up,down,ierr)
!Prepare for the passing of boundaries using MPI_PROC_NULL
if(left==-1) left=MPI_PROC_NULL
if(right==-1) right=MPI_PROC_NULL
if(down==-1) down=MPI_PROC_NULL
if(up==-1) up=MPI_PROC_NULL

if(rank==0)then
  print*,"nproc_x",dims(1)
  print*,"nproc_y",dims(2)
endif

!Split the grid on the processors or exit if the grid can't be split correctly
!-- make sure it always happens that nx/nproc_x = nx_glpbal etc.
nx = nx_global/nproc_x
ny = ny_global/nproc_y
IF (nx * nproc_x /= nx_global .OR. ny * nproc_y /= ny_global) THEN
  PRINT *,"Unable to evenly subdivide the grid on ",nproc," processors."
  PRINT *,"Please respecify."
  CALL MPI_ABORT(comm,ierr)
  CALL MPI_FINALIZE(ierr)
  STOP
ENDIF

end subroutine

!*********************************************************************************************

subroutine PASSINFO(nx,ny,nx_global,ny_global,nproc_x,nproc_y,rank,nextrank,comm, &
                    x,x_glob,y,y_glob,dx,dy,uold,IC_glob)
!This subroutine is used to distribute the sub grids to each of the slave
!processes. The array IC_glob is split into an array uold on each process.
IMPLICIT NONE
include 'mpif.h'
INTEGER :: ix,iy,nproc_x,nproc_y,rank,ierr,tag=100
INTEGER, DIMENSION(MPI_STATUS_SIZE) :: stat
INTEGER :: nx,ny,nx_global,ny_global
INTEGER :: coords(2),comm,nextrank
REAL :: x(0:nx+1),x_glob(nx_global),y(0:ny+1),y_glob(ny_global)
REAL :: IC_glob(nx_global,ny_global),uold(0:nx+1,0:ny+1)
REAL :: dx,dy


IF (rank == 0) THEN
  !On node 0 do all the sending
  DO ix=0,nproc_x-1
    DO iy=0,nproc_y-1
      !MPI_CART_RANK gets a process's rank from it's coordinates in the Cartesian topolgy.
      !You normally won't need it
      coords(1)=ix
      coords(2)=iy
      CALL MPI_CART_RANK(comm,coords,nextrank,ierr)
      IF (nextrank == 0) THEN
        !If the next rank is 0 (the one with the whole grid) then simply copy into place
        x(1:nx)=x_glob(ix * nx + 1 : (ix + 1) * nx)
        y(1:ny)=y_glob(iy * ny + 1 : (iy + 1) * ny)
        uold(1:nx,1:ny)=IC_glob(ix * nx + 1 : (ix + 1) * nx,iy * ny + 1 : (iy + 1) * ny)
         
      ELSE
        !Otherwise use MPI call
        CALL MPI_SEND(x_glob(ix*nx+1:(ix+1)*nx),nx,MPI_REAL,nextrank,tag,comm,ierr)
        CALL MPI_SEND(y_glob(iy*ny+1:(iy+1)*ny),ny,MPI_REAL,nextrank,tag,comm,ierr)
        CALL MPI_SEND(IC_glob(ix*nx+1:(ix+1)*nx,iy*ny+1:(iy+1)*ny),nx*ny,MPI_REAL,nextrank,tag,comm,ierr)
        CALL MPI_SEND(dx,1,MPI_REAL,nextrank,tag,comm,ierr)
        CALL MPI_SEND(dy,1,MPI_REAL,nextrank,tag,comm,ierr)
      ENDIF
    ENDDO
  ENDDO
ELSE
  !Every other process just wait for its section of the global array
  CALL MPI_RECV(x(1:nx),nx,MPI_REAL,0,tag,comm,stat,ierr)
  CALL MPI_RECV(y(1:ny),ny,MPI_REAL,0,tag,comm,stat,ierr)
  CALL MPI_RECV(uold(1:nx,1:ny),nx*ny,MPI_REAL,0,tag,comm,stat,ierr)
  CALL MPI_RECV(dx,1,MPI_REAL,0,tag,comm,stat,ierr)
  CALL MPI_RECV(dy,1,MPI_REAL,0,tag,comm,stat,ierr)
ENDIF

return
end subroutine

!********************************************************************************************************

subroutine COLLECTINFO(nx,ny,nx_global,ny_global,nproc_x,nproc_y,rank,nextrank,comm,uold,IC_glob)
!This subroutine collects all of the sub grids back together on the master
!process (i.e. the one with rank 0) and stores the final solution in the array
!IC_glob.
IMPLICIT NONE
include 'mpif.h'
INTEGER :: ix,iy,nproc_x,nproc_y,rank,ierr,tag=100
INTEGER, DIMENSION(MPI_STATUS_SIZE) :: stat
INTEGER :: nx,ny,nx_global,ny_global
INTEGER :: coords(2),comm,nextrank
REAL :: IC_glob(nx_global,ny_global),uold(0:nx+1,0:ny+1),usend(nx,ny)

usend(:,:)=uold(1:nx,1:ny)

IF (rank.ne.0) THEN
  !Every process send uold (for now) to processor 0
  CALL MPI_SEND(usend,nx*ny,MPI_REAL,0,tag,comm,ierr)
ELSE
  !On node 0 do all the receiving
  DO ix=0,nproc_x-1
    DO iy=0,nproc_y-1
      !MPI_CART_RANK gets a process's rank from it's coordinates in the Cartesian topolgy.
      coords(1)=ix
      coords(2)=iy
      CALL MPI_CART_RANK(comm,coords,nextrank,ierr)
      IF (nextrank == 0) THEN
        !If the next rank is 0 (the one with the whole grid) then simply copy into place
        IC_glob(ix*nx+1:(ix+1)*nx,iy*ny+1:(iy+1)*ny)=usend(1:nx,1:ny)
      ELSE
        !Otherwise use MPI call
        CALL MPI_RECV(IC_glob(ix*nx+1:(ix+1)*nx,iy*ny+1:(iy+1)*ny),nx*ny,MPI_REAL,nextrank,tag,comm,stat,ierr)
      ENDIF
    ENDDO
  ENDDO
ENDIF

return
end subroutine

!*********************************************************************************************


subroutine GETGLOBIC(nx_global,ny_global,dx,dy,x_glob,y_glob,IC_glob)
!This subroutine defines the mesh in x and y and the initial condition
implicit none
integer :: nx_global,ny_global,i,j
real :: IC_glob(nx_global,ny_global),dx,dy
real :: x_glob(nx_global),y_glob(ny_global)

do i=1,nx_global
  x_glob(i)=i*dx
  do j=1,ny_global
    y_glob(j)=j*dy
    IC_glob(i,j)=1.0*exp(-(1.0/1000.0)*((y_glob(j)-500.)**2+(x_glob(i)-500.)**2))
  end do
end do
end subroutine
!*********************************************************************************************
