c     program heat-2
c------------------------------------------------------------
c     Numerical solution of 1-D heat equation with 
c     explict scheme
c------------------------------------------------------------
      include "mpif.h"
      integer my_rank,nproc,ierr,n
      real*8 h,r,k,xl,u1(200),u2(200),x(200)
      real*8 xx(400),array(400)
      call init(nproc,my_rank)
c
c     define number of of computational points for each process
      n=2
c
c     total number of points: grid pts for each process + 2 bndry pts
      ntot=n*nproc+2

      xl=1.0

c     spacing h
      h=xl/dfloat(ntot-1)

c     r=k/h^2=0.25
      r=0.25
      k=r*h*h
      call init_con(nproc,my_rank,n,h,x,u1)
      do i=1,2
         if(my_rank.eq.0) then
         endif
         call comput(nproc,my_rank,n,r,u1,u2)
         call exchange(nproc,my_rank,n,u2)
         call comput(nproc,my_rank,n,r,u2,u1)
         call exchange(nproc,my_rank,n,u1)
         iroot=0
         call MPI_Gather(x(2),n,MPI_DOUBLE_PRECISION,xx,n,
     *     MPI_DOUBLE_PRECISION,iroot,MPI_COMM_WORLD,ierr)
         call MPI_Gather(u1(2),n,MPI_DOUBLE_PRECISION,array,n,
     *     MPI_DOUBLE_PRECISION,iroot,MPI_COMM_WORLD,ierr)
         if(my_rank.eq.0) then
            write(*,*) '***time level=',i*2,'***'
            do j=1,n*nproc
               write(*,105) j,xx(j),array(j)
 105           format(i5,2(f16.9,2x))
            enddo
         endif
      enddo
      call MPI_Finalize(ierr)
      stop
      end
c
c--------------------------------------------------------
c     Initial condition
c--------------------------------------------------------
      subroutine init_con(nproc,my_rank,n,h,x,u1)
      real*8 x(100),u1(100),h
      integer nproc,my_rank,n
      do i=1,n+1
         x(i)=0.0
      enddo
c     determine x coord
      do i=1,n+2
         x(i)=my_rank*dble(n)*h+dble(i-1)*h
      enddo
      write(*,*) 'my_rank=',my_rank, '  x',(x(i),i=1,n+2)
c     compute the initial condition
      pi=dacos(-1.d0)
      do i=1,n+2
         u1(i)=dsin(pi*x(i))
      enddo
      return
      end
c
c-----------------------------------------------------------
c     Exchangeing the first and the last u to the neighbors
c-----------------------------------------------------------
      subroutine exchange(nproc,my_rank,n,u)
      include "mpif.h"
      real*8 u(100)
      integer status(MPI_STATUS_SIZE)
      integer nproc,my_rank,itag,idest,isrc,nproc1,l_nbr,r_nbr
      nproc1=nproc-1
c
c     sending the first u to my_rank-1 for 0< my_rank < nproc-1
c
      if((my_rank-1).lt.0) then
         l_nbr=MPI_PROC_NULL
      else
         l_nbr=my_rank-1
      endif
      if((my_rank+1).gt.nproc1) then
         r_nbr=MPI_PROC_NULL
      else
         r_nbr=my_rank+1
      endif
      write(*,*) 'my_rank=',my_rank,'  l_nbr=',l_nbr,'  r_nbr=',r_nbr
      itag=20
c
c     send u(2) to left neighbor and receive u(n+2) from right neighbor
c
      call MPI_Sendrecv(u(2),1,MPI_DOUBLE_PRECISION,l_nbr,itag,
     *     u(n+2),1,MPI_DOUBLE_PRECISION,r_nbr,itag,
     *     MPI_COMM_WORLD,status,ierr)
c     
c     sending the last u to my_rank+1 for 0< my_rank < nproc-1
c     
      itag=10
c
c     send u(n+1) to right neighbor and receive u(1) from left neighbor
c
      call MPI_Sendrecv(u(n+1),1,MPI_DOUBLE_PRECISION,r_nbr,itag,
     *     u(1),1,MPI_DOUBLE_PRECISION,l_nbr,itag,
     *     MPI_COMM_WORLD,status,ierr)
      return
      end
c
c--------------------------------------------------------------
c     Computing u for the next time step
c--------------------------------------------------------------
      subroutine comput(nproc,my_rank,n,r,u1,u2)
      real*8 u1(100),u2(100),r
      integer nproc,my_rank,n
      if(my_rank.eq.0) then
         u2(1)=0.0
      else
         if(my_rank.eq.nproc-1) then
            u2(n+2)=0.0
         endif
      endif
c
c     excluding the ghost points at two ends
c
      do i=2,n+1
         u2(i)=r*(u1(i+1)-2.*u1(i)+u1(i-1))+u1(i)
      enddo
      return
      end
c
c-----------------------------------------------------------
c     Initializing MPI
c-----------------------------------------------------------
      subroutine init(nproc,my_rank)

      include "mpif.h"
c
c     initialize the MPI
c
      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