c     program heat-1d_x
c     1-D domain decomposition in x-direction
c
      include "mpif.h"
      parameter (MAX_X=400,MAX_Y=200)
      integer my_rank,nproc,ierr,n,stride,ntotx
      real*8 h,r,k,yl,u1(MAX_X,MAX_Y),u2(MAX_X,MAX_Y)
      real*8 xx(400),array(400),diff1,diffnorm,uu(400)
c     
c     INITIALIZE MPI AND DETERMINE BOTH INDIVIDUAL PROCESSOR #
c     AND THE TOTAL NUMBER OF PROCESSORS
c     
      call init(nproc,my_rank)
c
c     r=k/h^2=0.25
c
c     define number of computational points for each process in x-dir
c
      ny=2
c
c     total number of points: grid pts for each process + 2 bndry pts
c
      ntoty=ny*nproc+2
c 
c     let total number of grid pts in y-dir is the same as that in x-dir
c
      ntotx=ntoty
      yl=1.0
c     spacing h
      h=yl/dfloat(ntotx-1)
      call init_con(nproc,my_rank,ny,MAX_X,ntotx,h,u1)
      diffnorm=10.
      iter=0
      do while(dabs(diffnorm).gt.1.e-4)
         call comput(nproc,my_rank,ny,MAX_X,ntotx,h,u1,u2)
         call diff(nproc,my_rank,ny,MAX_X,ntotx,u1,u2,diff1)
         call MPI_Allreduce(diff1,diffnorm,1,MPI_DOUBLE_PRECISION,
     *        MPI_SUM,MPI_COMM_WORLD,ierr)
         iter=iter+1
         if(my_rank.eq.0) then
            write(*,*) 'diffnorm=',diffnorm,'   iter=',iter
         endif
         call exchange(nproc,my_rank,ny,MAX_X,ntotx,u2)
         call comput(nproc,my_rank,ny,MAX_X,ntotx,h,u2,u1)
         call exchange(nproc,my_rank,ny,MAX_X,ntotx,u1)
         call diff(nproc,my_rank,ny,MAX_X,ntotx,u1,u2,diff1)
         call MPI_Allreduce(diff1,diffnorm,1,MPI_DOUBLE_PRECISION,
     *        MPI_SUM,MPI_COMM_WORLD,ierr)
         iter=iter+1
         if(my_rank.eq.0) then
            write(*,*) 'diffnorm=',diffnorm,'   iter=',iter
         endif
      enddo
c
c     collocting u1 of each proc to global array
c
      iroot=0
      do i=2,ntotx-1
         do j=1,ny+2
            uu(j)=u1(i,j)
         enddo
         call MPI_Gather(uu(2),ny,MPI_DOUBLE_PRECISION,array,ny,
     *        MPI_DOUBLE_PRECISION,iroot,MPI_COMM_WORLD,ierr)
         if(my_rank.eq.0) then
            write(*,*) (array(j),j=1,ny*nproc)
         endif
      enddo
      call MPI_Finalize(ierr)
      stop
      end
c
c--------------------------------------------------------
c     Initial Guess * Boundary Condition
c--------------------------------------------------------
      subroutine init_con(nproc,my_rank,ny,MAX_X,ntotx,h,u1)
      real*8 u1(MAX_X,ny+2),h
      integer nproc,my_rank,ny
      do i=1,ntotx
         do j=1,ny+2
            u1(i,j)=0.0
         enddo
      enddo
c
c     bndry u(0,y)=u(x,1)=1, u(x,0)=u(1,y)=0
c
      do j=1,ny+2
         u1(1,j)=1.0
      enddo
      if(my_rank.eq.0) then
         do i=1,ntotx
            u1(i,1)=1.0
         enddo
      endif
      if(my_rank.eq.0) then
         do i=1,ntotx
            write(*,*) (u1(i,j),j=1,ny+2)
         enddo
      endif
      return
      end
c
c-----------------------------------------------------------
c     Exchangeing the first and the last u to the neighbors
c-----------------------------------------------------------
      subroutine exchange(nproc,my_rank,ny,MAX_X,ntotx,u)
      include "mpif.h"
      real*8 u(MAX_X,ny+2)
      integer status(MPI_STATUS_SIZE),stride
      integer nproc,my_rank,itag,idest,isrc,nproc1,b_nbr,t_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
         b_nbr=-1
      else
         b_nbr=my_rank-1
      endif
      if((my_rank+1).gt.nproc1) then
         t_nbr=-1
      else
         t_nbr=my_rank+1
      endif
c      write(*,*) 'my_rank=',my_rank,'  b_nbr=',b_nbr,'  t_nbr=',t_nbr
      itag=20
c
c     send u(2) to left neighbor and receive u(n+2) from right neighbor
c
      call MPI_Sendrecv(u(1,2),ntotx,MPI_DOUBLE_PRECISION,b_nbr,
     *     itag,u(1,ny+2),ntotx,MPI_DOUBLE_PRECISION,t_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(1,ny+1),ntotx,MPI_DOUBLE_PRECISION,t_nbr,
     *     itag,u(1,1),ntotx,MPI_DOUBLE_PRECISION,b_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,ny,MAX_X,ntotx,h,u1,u2)
      real*8 u1(MAX_X,ny+2),u2(MAX_X,ny+2),h
      real*8 f(400,200)
      integer nproc,my_rank
      do i=1,ntotx
         do j=1,ny+2
            u2(i,j)=u1(i,j)
         enddo
      enddo
      do i=2,ntotx-1
         do j=2,ny+1
            f(i,j)=0.0
            u2(i,j)=0.25*(u1(i-1,j)+u1(i,j+1)+u1(i,j-1)+u1(i+1,j)-
     *           h*h*f(i,j))
         enddo
      enddo
      if(my_rank.eq.0) then
         write(*,*) 'my_rank=',my_rank
         do i=1,ntotx
            write(*,*) (u2(i,j),j=1,ny+2)
         enddo
      endif
      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
c----------------------------------------------------------
      subroutine diff(nproc,my_rank,ny,MAX_X,ntotx,u1,u2,diff1)
      real*8 u1(MAX_X,ny+2),u2(MAX_X,ny+2),h,diff1
      diff1=0.0d0
      do i=1,ntotx
         do j=2,ny+1
            diff1=diff1+(u1(i,j)-u2(i,j))**2
         enddo
      enddo
      return
      end