c     domain-dcmp/heat/heat-2d_xy
c     2-D domain decomposition
c
      include "mpif.h"
      parameter (NDIM=2)
      parameter (MAX_X=200,MAX_Y=200)
      integer my_rank,nproc,ierr,stride1,stride2
      real*8 h1,h2,r1,r2,ua(MAX_X,MAX_Y),ub(MAX_X,MAX_Y),dt,t
      real*8 va(MAX_X,MAX_Y),vb(MAX_X,MAX_Y)
      integer comm2d,my_coord(2)
      integer l_nbr,r_nbr
      integer t_nbr,b_nbr,my_crank,my_coord(NDIM)
c
      call init(nproc,my_rank)
      call nbr2d(comm2d,my_crank,my_coord,l_nbr,r_nbr,t_nbr,b_nbr)
c
c     fine grid nx2=8, r2=dt/h^2=0.25
      r2=0.25
      nx2=8
      ny2=8
      ntotx2=nx2*2+2
      ntoty2=ny2*2+2
      xl=1.0
      h2=xl/(ntotx2-1)
      dt=r2*h2*h2

c     coarse grid nx1=4
      nx1=4
      ny1=4
      ntotx1=nx1*2+2
      ntoty1=ny1*2+2
      xl=1.0
      h1=xl/(ntotx1-1)
      r1=dt/(h1*h1)
c
c     define new type for message passing for coarse grid
c     ntoty number of block
c     1:     number of elements in each block
c     MAX_X:  number of elements between start of each block
c     MPI_DOUBLE_PRECISION: original type
      call MPI_Type_vector(ny1+2,1,MAX_X,
     *     MPI_DOUBLE_PRECISION,stride1,ierr)
      call MPI_type_commit(stride1,ierr)
c
c     new type for fine grid
      call MPI_Type_vector(ny2+2,1,MAX_X,
     *     MPI_DOUBLE_PRECISION,stride2,ierr)
      call MPI_type_commit(stride2,ierr)
c
      call init_cn(my_coord,my_crank,MAX_X,nx1,ny1,h1,ua)
      call init_cn(my_coord,my_crank,MAX_X,nx2,ny2,h2,va)
c
c     4 time levels
      do k=1,4
c
c        coarse grid
         t=dt*dble(k-1)
         call bndry_cn(my_coord,my_crank,MAX_X,nx1,ny1,ua,t)
         call comput(my_crank,MAX_X,nx1,ny1,r1,ua,ub)
         call exchange(comm2d,l_nbr,r_nbr,t_nbr,b_nbr,MAX_X,
     *        nx1,ny1,stride1,ub)
         if(my_rank.eq.3) then
            write(*,*) 'final'
            do j=1,ny1+2
               write(*,*) 'k=',k,' j=',j,'ub',(ub(i,j),i=1,nx1+2)
            enddo
         endif
         do i=1,nx1+2
            do j=1,ny1+2
               ua(i,j)=ub(i,j)
            enddo
         enddo
c
c        fine grid
         call bndry_cn(my_coord,my_crank,MAX_X,nx2,ny2,va,t)
         call comput(my_crank,MAX_X,nx2,ny2,r2,va,vb)
         call exchange(comm2d,l_nbr,r_nbr,t_nbr,b_nbr,MAX_X,
     *        nx2,ny2,stride2,vb)
         if(my_rank.eq.3) then
            write(*,*) 'final'
            do j=2,ny2+1
               write(*,*) 'k=',k,' j=',j,'vb',(vb(i,j),i=2,nx2+1)
            enddo
         endif
         do i=1,nx2+2
            do j=1,ny2+2
               va(i,j)=vb(i,j)
            enddo
         enddo
      enddo
      call MPI_Finalize(ierr)
      stop
      end
c
c-------------------------------------------------------
c     Cartisian Grid for Processes
c------------------------------------------------------
      subroutine nbr2d(comm2d,my_crank,my_coord,
     *     l_nbr,r_nbr,t_nbr,b_nbr)
      include "mpif.h"
      parameter (NDIM=2)
      integer idim(2),comm2d,my_crank,l_nbr,r_nbr,t_nbr,b_nbr
      integer sideways,updown,right,up,my_coord(NDIM)
      logical periods(2),reorder
      idim(1)=2
      idim(2)=2
      periods(1)=.false.
      periods(2)=.false.
      reorder=.true.
      sideways=1
      updown=0
      right=1
      up=1
      call MPI_Cart_Create(MPI_COMM_WORLD,NDIM,idim,
     *     periods,reorder,comm2d, ierr)
      call MPI_Comm_Rank(comm2d,my_crank,ierr)
      call MPI_Cart_Coords(comm2d,my_crank,NDIM,my_coord,ierr)
      call MPI_Cart_Shift(comm2d,sideways,right,l_nbr,r_nbr,ierr)
      call MPI_Cart_Shift(comm2d,updown,up,t_nbr,b_nbr,ierr)
      write(*,*) 'my_crank=',my_crank,
     *     '  coor=(',my_coord(1),',',my_coord(2),' )'
      write(*,*) 'my_crank=',my_crank,' l_nbr=',l_nbr,
     *     ' r_nbr=',r_nbr,
     * 't_nbr=',t_nbr,'  b_nbr=',b_nbr
      return
      end
c
c--------------------------------------------------------
c     Initial Condition
c--------------------------------------------------------
      subroutine init_cn(my_coord,my_rank,MAX_X,nx,ny,h,u1)
      include "mpif.h"
      real*8 u1(MAX_X,1),h
      real*8 x(100),y(100),pi
      integer my_coord(2),nx,ny
      pi=dacos(-1.d0)
      do i=1,nx+2
         x(i)=dble(i-1)*h
      enddo
      do j=1,ny+2
         y(j)=dble(j-1)*h
      enddo
      if(my_coord(1).eq.0) then
         dh=h*dble(ny)
         do j=1,ny+2
            y(j)=y(j)+dh
         enddo
      endif
      if(my_coord(2).eq.1) then
         dh=h*dble(nx)
         do i=1,nx+2
            x(i)=dh+x(i)
         enddo
      endif
c
c     initial condition
      do i=1,nx+2
         do j=1,ny+2
            u1(i,j)=dsin(pi*x(i))*dsin(pi*y(j))
         enddo
      enddo
      return
      end
c
c-----------------------------------------------------------
c     Boundary Condition
c-----------------------------------------------------------
      subroutine bndry_cn(my_coord,my_rank,MAX_X,nx,ny,u,t)
      integer my_coord(2),nx
      real*8 u(MAX_X,1),t
c
c     bndry u(x,0)=sin(t)
      if(my_coord(1).eq.1) then
         do i=1,nx+2
            u(i,1)=dsin(t)
         enddo
      endif

c     bndry u(x,1)=sin(t)
      if(my_coord(1).eq.0) then
         do i=1,nx+2
            u(i,ny+2)=dsin(t)
         enddo
      endif
c
c     bndry u(0,y)=sin(t)
      if(my_coord(2).eq.0) then
         do j=1,ny+2
            u(1,j)=dsin(t)
         enddo
      endif
c
c     bndry u(1,y)=sin(t)
      if(my_coord(2).eq.1) then
         do j=1,ny+2
            u(nx+2,j)=dsin(t)
         enddo
      endif
      return
      end
c
c-----------------------------------------------------------
c     Exchangeing the first and the last u to the neighbors
c-----------------------------------------------------------
      subroutine exchange(comm2d,l_nbr,r_nbr,t_nbr,b_nbr,MAX_X,
     *     nx,ny,stride,u)
      include "mpif.h"
      integer comm2d
      real*8 u(MAX_X,1)
      integer status(MPI_STATUS_SIZE),stride
      integer l_nbr,r_nbr,t_nbr,b_nbr,itag

      itag=10
c     send u(2,j) to left neighbor and receive u(n+2,j) from right neighbor
c
      call MPI_Sendrecv(u(2,1),1,stride,l_nbr,itag,
     *     u(nx+2,1),1,stride,r_nbr,itag,
     *     comm2d,status,ierr)
c     
c     sending the last u to my_rank+1 for 0< my_rank < nproc-1
c     
      itag=20
c     send u(n+1,j) to right neighbor and receive u(1,j) from left neighbor
c
      call MPI_Sendrecv(u(nx+1,1),1,stride,r_nbr,itag,
     *     u(1,1),1,stride,l_nbr,itag,
     *     comm2d,status,ierr)
c
c     send u(i,2) to bottom neighbor and receive u(i,n+2) from top neighbor
c
      itag=30
      call MPI_Sendrecv(u(1,2),nx+2,MPI_DOUBLE_PRECISION,b_nbr,
     *     itag,u(1,ny+2),nx+2,MPI_DOUBLE_PRECISION,t_nbr,itag,
     *     comm2d,status,ierr)
c
c     sending the last u to my_rank+1 for 0< my_rank < nproc-1
c
      itag=40
c     send u(i,n+1) to top neighbor and receive u(i,1) from bottom neighbor
c
      call MPI_Sendrecv(u(1,ny+1),nx+2,MPI_DOUBLE_PRECISION,t_nbr,
     *     itag,u(1,1),nx+2,MPI_DOUBLE_PRECISION,b_nbr,itag,
     *     comm2d,status,ierr)
      return
      end
c
c--------------------------------------------------------------
c     Computing u for the next time step
c--------------------------------------------------------------
      subroutine comput(my_rank,MAX_X,nx,ny,r,u1,u2)
      real*8 u1(MAX_X,1),u2(MAX_X,1),r
      do i=1,nx+2
         do j=1,ny+2
            u2(i,j)=u1(i,j)
         enddo
      enddo
      do i=2,nx+1
         do j=2,ny+1
            u2(i,j)=r*(u1(i-1,j)-2.*u1(i,j)+u1(i+1,j)+
     *           u1(i,j-1)-2.*u1(i,j)+u1(i,j+1))+u1(i,j)
         enddo
      enddo
      return
      end
c
      subroutine gather(k,nproc,my_rank,MAX_X,nx,ny,u)
      integer k,iroot,nproc,my_rank,nx,ntoty
      real*8 u(MAX_X,1),uu(200),array(400)
      include 'mpif.h'
c     collocting u of each proc to global array
      iroot=0
      do j=2,ny+1
         do i=1,nx+2
            uu(i)=u(i,j)
         enddo
         call MPI_Gather(uu(2),nx,MPI_DOUBLE_PRECISION,array,nx,
     *        MPI_DOUBLE_PRECISION,iroot,MPI_COMM_WORLD,ierr)
         if(my_rank.eq.0) then
            write(*,*) 'k=',k,' j=',j,' ',(array(i),i=1,nx*nproc)
         endif
      enddo
      return
      end
c
c-----------------------------------------------------------
c     Initializing MPI
c-----------------------------------------------------------
      subroutine init(nproc,my_rank)
      include "mpif.h"
      integer my_rank,nproc
      character*30 nodename
      integer nchar
      call MPI_init(ierr)
      call MPI_Comm_rank(MPI_COMM_WORLD,my_rank,ierr)
      call MPI_Comm_size(MPI_COMM_WORLD,nproc,ierr)
      call MPI_Get_processor_name(nodename,nchar,ierr)
      write(*,*) 'my_rank=',my_rank,' processor=',nodename
      return
      end