c     program heat-1d_x
c     1-D domain decomposition in x-direction
c
      include "mpif.h"
      parameter (MAX_X=200,MAX_Y=400)
      integer my_rank,nproc,ierr,n,stride
      real*8 h,r,k,xl,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
      nx=2
c
c     total number of points: grid pts for each process + 2 bndry pts
c
      ntotx=nx*nproc+2
c 
c     let total number of grid pts in y-dir is the same as that in x-dir
c
      ntoty=ntotx
      xl=1.0
c     spacing h
      h=xl/dfloat(ntotx-1)
      call init_con(nproc,my_rank,nx,MAX_X,ntoty,h,u1)
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: 
       call MPI_Type_vector(ntoty,1,MAX_X,
     *     MPI_DOUBLE_PRECISION,stride,ierr)
      call MPI_type_commit(stride,ierr)
      diffnorm=10.
      iter=0
      do while(dabs(diffnorm).gt.1.e-4)
         call comput(nproc,my_rank,nx,MAX_X,ntoty,h,u1,u2)
         call exchange(nproc,my_rank,nx,MAX_X,ntoty,stride,u2)
         call comput(nproc,my_rank,nx,MAX_X,ntoty,h,u2,u1)
         call exchange(nproc,my_rank,nx,MAX_X,ntoty,stride,u1)
         call diff(nproc,my_rank,nx,MAX_X,ntoty,u1,u2,diff1)
         call MPI_Allreduce(diff1,diffnorm,1,MPI_DOUBLE_PRECISION,
     *        MPI_SUM,MPI_COMM_WORLD,ierr)
         iter=iter+2
         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 j=2,ntoty-1
         do i=1,nx+2
            uu(i)=u1(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(*,*) (array(i),i=1,nx*nproc)
         endif
      enddo
      call MPI_Finalize(ierr)
      stop
      end
c
c--------------------------------------------------------
c     Initial Guess * Boundary Condition
c--------------------------------------------------------
      subroutine init_con(nproc,my_rank,nx,MAX_X,ntoty,h,u1)
      real*8 u1(MAX_X,ntoty),h
      integer nproc,my_rank,nx
      do i=1,nx+2
         do j=1,ntoty
            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
      if(my_rank.eq.0) then
         do j=1,ntoty
            u1(1,j)=1.0
         enddo
      endif
      if(my_rank.eq.nproc-1) then
         do i=j,ntoty
            u1(nx+2,j)=0.0
         enddo
      endif
      do i=1,nx+2
         u1(i,1)=1.0
      enddo
c      if(my_rank.eq.2) then
c         do i=1,nx+2
c            write(*,*) (u1(i,j),j=1,ntoty)
c         enddo
c      endif
      return
      end
c
c-----------------------------------------------------------
c     Exchangeing the first and the last u to the neighbors
c-----------------------------------------------------------
      subroutine exchange(nproc,my_rank,nx,MAX_X,ntoty,
     *     stride,u)
      include "mpif.h"
      real*8 u(MAX_X,ntoty)
      integer status(MPI_STATUS_SIZE),stride
      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=-1
      else
         l_nbr=my_rank-1
      endif
      if((my_rank+1).gt.nproc1) then
         r_nbr=-1
      else
         r_nbr=my_rank+1
      endif
c      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),1,stride,l_nbr,itag,
     *     u(nx+2,1),1,stride,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(nx+1,1),1,stride,r_nbr,itag,
     *     u(1,1),1,stride,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,nx,MAX_X,ntoty,h,u1,u2)
      real*8 u1(MAX_X,ntoty),u2(MAX_X,ntoty),h
      real*8 f(200,400)
      integer nproc,my_rank
      do i=1,nx+2
         do j=1,ntoty
            u2(i,j)=u1(i,j)
         enddo
      enddo
      do i=2,nx+1
         do j=2,ntoty-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
c      if(my_rank.eq.3) then
c         write(*,*) 'my_rank=',my_rank
c         do i=2,nx+1
c            write(*,*) (u2(i,j),j=1,ntoty)
c         enddo
c      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,nx,MAX_X,ntoty,u1,u2,diff1)
      real*8 u1(MAX_X,ntoty),u2(MAX_X,ntoty),h,diff1
      diff1=0.0d0
      do i=2,nx+1
         do j=1,ntoty
            diff1=diff1+(u1(i,j)-u2(i,j))**2
         enddo
      enddo
      return
      end