program heat2D
use mpi
implicit none

 integer :: niter,i,j,t,rank, ndim, nbproc, ierror,MPI_COMM_CART
 real(8), allocatable, dimension(:,:) :: U0,U1
 real(8) :: alpha,totaltime,Dt,Stab,DtAlpha
 integer, dimension(1:2) :: Nx, Nxl, nbproc_axe, cart_position
 real(8), dimension(1:2) :: Lx,Dx,DxDx,InvDxDx
 character(len=4) :: charbuff1,charbuff2
 integer, dimension(-1:1,-1:1) :: cart_neigh
 logical, dimension(1:2) :: cart_boundaries

 cart_boundaries(1) = .false.
 cart_boundaries(2) = .true.
 niter = 8000
 ndim = 2
 Nx(:) = 800

 nbproc_axe(1) = 2
 nbproc_axe(2) = 2

 ! Assuming we can divide number of points by number of processes per axes
 ! Nxl is the number of grid points for eaxh process
 Nxl(:) = Nx(:)/nbproc_axe(:)

 alpha = 0.4d0
 totaltime = 0.005d0
 Dt = totaltime/((niter-1)*1.0d0)
 Lx(:) = 1.0d0
 Dx(:) = Lx(:)/((Nx(:)-1)*1.0d0)
 InvDxDx(:) = 1.0d0/(Dx(:)*Dx(:))
 DxDx(:) = Dx(:)*Dx(:)
 Stab = alpha*Dt*(InvDxDx(1)) + alpha*Dt*(InvDxDx(2))
 DtAlpha = Dt*alpha
 print *, "Stability factor (need < 0.5): ",stab
 if(stab > 0.5d0) stop
 print *,"Calculating, please wait..."

 allocate(U0(0:Nxl(1)+1,0:Nxl(2)+1),U1(0:Nxl(1)+1,0:Nxl(2)+1))


 call MPI_INIT( ierror )
 call MPI_COMM_SIZE( MPI_COMM_WORLD , nbproc , ierror )

 call MPI_CART_CREATE( MPI_COMM_WORLD , ndim , nbproc_axe(1:ndim) , &
 & cart_boundaries(1:ndim) , .true. , MPI_COMM_CART , ierror )

 call MPI_COMM_RANK( MPI_COMM_CART , rank , ierror )

 call MPI_CART_COORDS( MPI_COMM_CART , rank , ndim , cart_position(1:ndim) , ierror )

 call MPI_CART_SHIFT (MPI_COMM_CART, 0, 1, cart_neigh(-1,0), cart_neigh(+1,0), ierror)
 call MPI_CART_SHIFT (MPI_COMM_CART, 1, 1, cart_neigh(0,-1), cart_neigh(0,+1), ierror)

 U0(:,:) = 1.0d0
 if(rank==0) U0(100:150,200:250) = -10.0d0
 if(rank==2) U0(300:350,150:200) = +12.0d0
 U1(:,:) = U0(:,:)

 write(charbuff1,'(I4.4)') cart_position(1)
 write(charbuff2,'(I4.4)') cart_position(2)
 open(10,file=trim(adjustl('IN'//'.'//trim(adjustl(charbuff1))//'.'//trim(adjustl(charbuff2))//'.dat')))
 do j=1,Nxl(2),10 ! do not write all data
    do i=1,Nxl(1),10
       write(10,*) i+(cart_position(1)*Nxl(1)),j+(cart_position(2)*Nxl(2)),U0(i,j)
    end do
 end do
 close(10)

 do t=1,niter

    ! Recv left, send right
    call MPI_SENDRECV ( U0(Nxl(1),1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(1,0) , 7777 , &
    & U0(0,1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(-1,0) , 7777 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )
    ! Recv right, send left
    call MPI_SENDRECV ( U0(1,1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(-1,0) , 7778 , &
    & U0(Nxl(1)+1,1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(1,0) , 7778 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )

    ! Recv down, send up
    call MPI_SENDRECV ( U0(1:Nxl(1),Nxl(2)) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,1) , 7779 , &
    & U0(1:Nxl(1),0) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,-1) , 7779 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )
    ! Recv up, send down
    call MPI_SENDRECV ( U0(1:Nxl(1),1) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,-1) , 7780 , &
    & U0(1:Nxl(1),Nxl(2)+1) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,1) , 7780 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )

    do j=1,Nxl(2)
       do i=1,Nxl(1)

          U1(i,j) = U0(i,j) + DtAlpha* &
          & ( (U0(i+1,j)-2.0*U0(i,j)+U0(i-1,j))*InvDxDx(1) + &
          &   (U0(i,j+1)-2.0*U0(i,j)+U0(i,j-1))*InvDxDx(2) )

       end do
    end do

    ! Recv left, send right
    call MPI_SENDRECV ( U1(Nxl(1),1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(1,0) , 7781 , &
    & U1(0,1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(-1,0) , 7781 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )
    ! Recv right, send left
    call MPI_SENDRECV ( U1(1,1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(-1,0) , 7782 , &
    & U1(Nxl(1)+1,1:Nxl(2)) , Nxl(2) , MPI_DOUBLE_PRECISION , cart_neigh(1,0) , 7782 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )

    ! Recv down, send up
    call MPI_SENDRECV ( U1(1:Nxl(1),Nxl(2)) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,1) , 7783 , &
    & U1(1:Nxl(1),0) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,-1) , 7783 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )
    ! Recv up, send down
    call MPI_SENDRECV ( U1(1:Nxl(1),1) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,-1) , 7784 , &
    & U1(1:Nxl(1),Nxl(2)+1) , Nxl(1) , MPI_DOUBLE_PRECISION , cart_neigh(0,1) , 7784 , MPI_COMM_CART , MPI_STATUS_IGNORE , ierror )

    do j=1,Nxl(2)
       do i=1,Nxl(1)

          U0(i,j) = U1(i,j) + DtAlpha* &
          & ( (U1(i+1,j)-2.0*U1(i,j)+U1(i-1,j))*InvDxDx(1) + &
          &   (U1(i,j+1)-2.0*U1(i,j)+U1(i,j-1))*InvDxDx(2) )

       end do
    end do

 end do

 write(charbuff1,'(I4.4)') cart_position(1)
 write(charbuff2,'(I4.4)') cart_position(2)
 open(10,file=trim(adjustl('OUT'//'.'//trim(adjustl(charbuff1))//'.'//trim(adjustl(charbuff2))//'.dat')))
 do j=1,Nxl(2),10 ! do not write all data
    do i=1,Nxl(1),10
       write(10,*) i+(cart_position(1)*Nxl(1)),j+(cart_position(2)*Nxl(2)),U0(i,j)
    end do
 end do
 close(10)

 call MPI_FINALIZE ( ierror ) ! Close MPI

 deallocate(U0,U1)

end program heat2D