c
c-----------------------------------------------------------------------
c     Conjugate Gradient Poisson Solver taken from MOM2.1
c
c     author:      Charles Goldberg        e-mail: chg@gfdl.gov
c
c     based on the preconditioned conjugate gradient algorithm given in:
c
c     A Reformulation and Implementation of the Bryan-Cox-Semtner
c     Ocean Model on the Connection Machine
c     J.K. Dukowicz, R.D. Smith, and R.C. Malone
c     Journal of Atmospheric and Oceanic Technology
c     Vol 10. No. 2 April 1993
c
c     Note: needs to be initialized by init_congrad
c           to calculate island integral paths.
c
c                      C.Eden Feb. 2002 (mailto ceden@phys.ocean.dal.ca)
c-----------------------------------------------------------------------
c
      module congrad_module
      implicit none
      integer :: imt,jmt,imain
      integer :: mnisle,maxipp,nisle
      logical, allocatable :: imask(:)
      integer, allocatable :: iperm(:),jperm(:)
      integer, allocatable :: iofs(:),nippts(:)
      integer, allocatable :: map(:,:)
      logical :: cyclic
      end module congrad_module


      subroutine init_congrad(imt_in,jmt_in,tmask,
     & cyclic_in, enable_obc_north,enable_obc_south,
     & enable_obc_west,enable_obc_east) 
      use island_module
      use congrad_module
      implicit none
      integer :: imt_in,jmt_in,kmt(imt_in,jmt_in),isle
      real(kind=8) :: tmask(imt_in,jmt_in)
      logical :: cyclic_in,enable_obc_north,enable_obc_south
      logical :: enable_obc_west,enable_obc_east

      print*,' initializing the poisson solver'

      imt=imt_in; jmt=jmt_in; cyclic = cyclic_in

      mnisle=10
      maxipp =100000
      allocate( imask(-mnisle:mnisle) )
      allocate( iperm(maxipp),jperm(maxipp) )
      allocate( iofs(mnisle),nippts(mnisle) )
      allocate( map(imt,jmt) )

      kmt=tmask
      call isleperim (kmt, map, iperm, jperm, iofs, nippts, nisle,
     & imt, jmt, mnisle, maxipp,0,cyclic, 
     & enable_obc_north,enable_obc_south,
     & enable_obc_west,enable_obc_east,.false.)  ! chamge kmt mask

      do isle=-mnisle,mnisle
        if (isle .ge. 0 .and. isle .le. nisle) then
          imask(isle) = .true.
        else
          imask(isle) = .false.
        end if
      end do
      print*,' done, here comes the land/sea mask:'
      call showmap(map,imt,jmt)
      imain=1
      end subroutine init_congrad


      subroutine set_cyclic (v)
      use congrad_module
      implicit none
      real (kind=8) :: v(imt,jmt)
      if (cyclic) then
       v(1,:)   = v(imt-1,:)
       v(imt,:) = v(2,:)
      endif
      end subroutine set_cyclic


      subroutine congrad(cf,dpsi,forc,res,
     &    epsilon,iterations,estimated_error,islands)
c     solve   A * dpsi = forc for dpsi
      use congrad_module
      implicit none
      real(kind=8), dimension(imt,jmt,-1:1,-1:1) :: cf
      real(kind=8), dimension(imt,jmt) :: dpsi,res,forc
      real(kind=8), dimension(imt,jmt) :: Z,Zres,As,s

      logical converged,diverged, islands
      integer iterations,max_iterations,k,i,j
      real(kind=8) epsilon,estimated_error,convergence_rate
      real(kind=8) Zresmax,betakm1,betak,betak_min,smax
      real(kind=8) alpha,step,betaquot,s_dot_as,step1,cfactor,dpsi1

      Z=0.; Zres=0.; As=0.; s=0.

      max_iterations=max(imt,jmt)*50

      call set_cyclic(dpsi)
c-----------------------------------------------------------------------
c     make approximate inverse operator Z (always even symmetry)
c-----------------------------------------------------------------------
c
      call make_inv (cf, Z, islands )
      call set_cyclic(Z)
c
c-----------------------------------------------------------------------
c     res(0)  = forc - A * dpsi(0)
c     impose cyclic and/or symmetry conditions on res(i,j)
c-----------------------------------------------------------------------
c
      call op9_vec(cf, dpsi, res)
      res = forc - res
      call set_cyclic(res)
c
c-----------------------------------------------------------------------
c     Zres(k-1) = Z * res(k-1)
c     see if guess is a solution, bail out to avoid division by zero
c-----------------------------------------------------------------------
c
      k = 0
      diverged=.false.

      call inv_op(Z, res, Zres, islands)
      call set_cyclic(Zres)
c
c
      Zresmax = maxval(abs(Zres))
c
c       Assume convergence rate of 0.99 to extrapolate error
c
      if (100.0 * Zresmax .lt. epsilon) then
	  estimated_error = 100.0 * Zresmax 
	  goto 101
      endif
c
c-----------------------------------------------------------------------
c     beta(0) = 1
c     s(0)    = zerovector()
c-----------------------------------------------------------------------
c
      betakm1 = 1.0
      s=0.
c
c-----------------------------------------------------------------------
c     begin iteration loop
c-----------------------------------------------------------------------
c
      do k = 1,max_iterations
c
c-----------------------------------------------------------------------
c       Zres(k-1) = Z * res(k-1)
c-----------------------------------------------------------------------
c
        call inv_op( Z, res, Zres, islands)
        call set_cyclic(Zres)

c-----------------------------------------------------------------------
c       beta(k)   = res(k-1) * Zres(k-1)
c-----------------------------------------------------------------------
c
        call dot2(betak,Zres, res)

        if (k .eq. 1) then
          betak_min = abs(betak)
	elseif (k .gt. 2) then
          betak_min = min(betak_min, abs(betak))
          if (abs(betak) .gt. 100.0*betak_min) then
	    write (6,'(/a/a/)')
     &      'WARNING: conjugate gradient terminated because correction'
     &,     '         steps are diverging. Probable cause...roundoff'
            print*,'k=',k
            smax = maxval(abs(s))
            step = abs(alpha) * smax
	    estimated_error=step*convergence_rate/(1.0-convergence_rate)
            diverged=.true.
	    go to 101
	  endif
	endif
c
c-----------------------------------------------------------------------
c       s(k)      = Zres(k-1) + (beta(k)/beta(k-1)) * s(k-1)
c-----------------------------------------------------------------------
c
        betaquot = betak/betakm1
        s = Zres + betaquot * s
c
c-----------------------------------------------------------------------
c       As(k)     = A * s(k)
c-----------------------------------------------------------------------
c

        call op9_vec(cf, s, As)
        call set_cyclic(As)
c
c-----------------------------------------------------------------------
c       If s=0 then the division for alpha(k) gives a float exception.
c       Assume convergence rate of 0.99 to extrapolate error.
c       Also assume alpha(k) ~ 1.
c-----------------------------------------------------------------------
c
        call dot2(s_dot_As ,s, As)

        if ((abs(s_dot_As) .lt. abs(betak)*1.e-10).or.
     &              ( s_dot_as == 0.)) then
          smax = maxval(abs(s))
	  estimated_error = 100.0 * smax 
	  goto 101
	endif
c
c-----------------------------------------------------------------------
c       alpha(k)  = beta(k) / (s(k) * As(k))
c-----------------------------------------------------------------------
c
        alpha = betak / s_dot_As
c
c-----------------------------------------------------------------------
c       update values:
c       dpsi(k)   = dpsi(k-1) + alpha(k) * s(k)
c       res(k)    = res(k-1) - alpha(k) * As(k)
c-----------------------------------------------------------------------
c
        dpsi  = dpsi + alpha * s
        res   = res  - alpha * As


        if (islands) call avg_dist (res)
        call set_cyclic(res)
c
        smax = maxval(abs(s))
c
c-----------------------------------------------------------------------
c       test for convergence
c       if (estimated_error) < epsilon) exit
c-----------------------------------------------------------------------
c
        step = abs(alpha) * smax
c        print*,'step=',step,' step1=',step1
        if (k .eq. 1) then
          step1 = step
          estimated_error = step
          if (step .lt. epsilon) goto 101
        else if (step .lt. epsilon) then
          if (step == 0) then
            cfactor=-39.
          else  
           cfactor = log(step/step1)
          endif
          convergence_rate = exp(cfactor/(k-1))
	  estimated_error = step*convergence_rate/(1.0-convergence_rate)
          if (estimated_error .lt. epsilon) goto 101
        end if
c
        betakm1 = betak
c
      end do
c
c-----------------------------------------------------------------------
c     end of iteration loop
c-----------------------------------------------------------------------
c
  101 continue
      if ((k .gt. max_iterations).or.(diverged)) then
        cfactor = log(step/step1)
        convergence_rate = exp(cfactor/(k-1))
	estimated_error = step*convergence_rate/(1.0-convergence_rate)
        converged = .false.
      else
        converged = .true.
      end if

      iterations = k
c
c-----------------------------------------------------------------------
c     return the last increment of dpsi in the argument res
c-----------------------------------------------------------------------
c
      if (iterations .eq. 0) then
        res = Zres
      else
        res = alpha * s
      endif

      if (islands) then
       dpsi1 = dpsi(iperm(iofs(imain)+1), jperm(iofs(imain)+1))
       do i=1,imt
        do j=1,jmt
          if (map(i,j) .le. 0) then
            dpsi(i,j) = dpsi(i,j) - dpsi1
          end if
        end do
       end do
      endif

      end subroutine congrad


c
c=======================================================================
c
c     M A T R I X   M O D U L E   F O R   C O N G R A D
c
c=======================================================================
c

      subroutine dot2(sum,a,b)
      use congrad_module
      implicit none
      real(kind=8) :: a(imt,jmt),b(imt,jmt),sum
      integer i,j
c
      sum=0.
      do j=1,jmt
       do i=2,imt-1
         sum=sum+a(i,j)*b(i,j)
       enddo
      enddo
      end subroutine dot2

      subroutine op9_vec(cf, dpsi, res)
c
c                       res = A * dpsi
c
c     this subroutine does not collect the terms of the true value
c     of res(isle) = sum (res(i,j)).  the contributions to the sum
c     remain distributed among the T cells (i,j) that form the
c     ocean perimeter of land_mass(isle).
c
c     at present, borders are not computed [i=1 or imt] [j=1 or jmt]
c
      use congrad_module
      implicit none
      real(kind=8) cf(imt,jmt,-1:1,-1:1), dpsi(imt,jmt),res(imt,jmt)
      integer i,j
c
      do j=2,jmt-1
        do i=2,imt-1
          res(i,j) =
     &               cf(i,j,-1,-1) * dpsi(i-1,j-1) +
     &               cf(i,j, 0,-1) * dpsi(i  ,j-1) +
     &               cf(i,j, 1,-1) * dpsi(i+1,j-1)
        end do
        do i=2,imt-1
          res(i,j) = res(i,j) +
     &               cf(i,j,-1, 0) * dpsi(i-1,j  ) +
     &               cf(i,j, 0, 0) * dpsi(i  ,j  ) +
     &               cf(i,j, 1, 0) * dpsi(i+1,j  )
        end do
        do i=2,imt-1
          res(i,j) = res(i,j) +
     &               cf(i,j,-1, 1) * dpsi(i-1,j+1) +
     &               cf(i,j, 0, 1) * dpsi(i  ,j+1) +
     &               cf(i,j, 1, 1) * dpsi(i+1,j+1)
        end do
      end do
      end subroutine op9_vec


      subroutine inv_op( Z, res, Zres, islands)
c
c     apply and approximate inverse Z or the operator A
c
c     res is type(res_type), i.e., perimeter values res(isle)
c         are the sum of the distributed contributions res(i,j)
c     Zres is type(dpsi_type), i.e., perimeter values Zres(isle)
c         must be replicated at each perimeter point Zres(i,j)
c
c     borders  of Zres [i=1 or imt] [j=1 or jmt] must be defined
c     and must satisfy cyclic and/or symmetry, if defined.
c
c     currently, Z is diagonal:  Z(ij) = 1/A(ij)
c     and is stored in type(dpsi_type) format, i.e., Z(isle) is
c     replicated and stored in each Z(i,j) in the perimeter of
c     land_mass(isle).
c
      use congrad_module
      implicit none
      real(kind=8) Z(imt,jmt), res(imt,jmt), Zres(imt,jmt)
      logical islands
c
      Zres = Z * res
c
c     sum contributions to Zres(isle)
c     distribute Zres(isle) to all perimeter points
      if (islands) call sum_dist (Zres)
c
      end subroutine inv_op



      subroutine make_inv (cf, Z, islands)
c
c     construct an approximate inverse Z to A
c
c     Z will be diagonal:  Z(ij) = 1/A(ij)
c     and values for ocean perimeter entries Z(isle) will be replicated
c     at all T cells Z(i,j) in the ocean perimeter of land_mass(isle).
c
c     T cells (i,j) for which there is no diagonal coefficient
c     i.e., A(ij)=A(i,j)=0, are masked off by assigning Z(i,j)=0.
c     there are effectively no equations and no variables dpsi(i,j)
c     at these points.
c
c
      use congrad_module
      implicit none
      real (kind=8)cf(imt,jmt,-1:1,-1:1), Z(imt,jmt)
      logical :: islands
      integer i,j,isle,n
c
c     copy diagonal coefficients of A to Z
c
      Z = cf(:,:,0,0)
c
c     for each land_mass(isle),
c     sum the contributions to cfdiag(isle)=A(isle,isle)
c     now stored in Z(i,j) at ocean perimeter T cells and replicate
c     the sum in all Z(i,j) for which (i,j) is in ocean perimeter
c     of land_mass(isle).
c
      if (islands) call sum_dist (Z)
c
c     now invert Z
c
      do j=1,jmt
        do i=1,imt
          if (Z(i,j) .ne. 0.0) then
            Z(i,j) = 1./Z(i,j)
          else
            Z(i,j) = 0.0
          end if
        end do
      end do
c
c     make inverse zero on island perimeters that are not integrated
c
      if (islands) then
       do isle=1,nisle
        if (.not. imask(isle)) then
          do n=1,nippts(isle)
            i = iperm(iofs(isle)+n)
            j = jperm(iofs(isle)+n)
            Z(i,j) = 0.0
          end do
        end if
       end do
      endif
c
      end subroutine make_inv



      subroutine sum_dist (Zres)
c
c     sum contributions to Zres(isle)
c     distribute Zres(isle) to all perimeter points
c
c     this subroutine converts a type(res_type) vector with
c     distributed contributions to perimeter values
c        Zres(isle) = sum (Zres(i,j))
c     into a type (dpsi_type) vector with replicated values
c     for land_mass perimeters
c        Zres(isle) = Zres(i,j)
c     for all (i,j) in the ocean perimeter of land_mass(isle).
c
      use congrad_module
      implicit none
      real(kind=8) Zres(imt,jmt), Zresisle(mnisle)
      integer i,j,isle,n
c
c     sum contributions to Zres(isle)
c
         do isle=1,nisle
           if (imask(isle)) then
             Zresisle(isle) = 0.0
             do n=1,nippts(isle)
               i = iperm(iofs(isle)+n)
               j = jperm(iofs(isle)+n)
               Zresisle(isle) = Zresisle(isle) + Zres(i,j)
             end do
           end if
         end do
c
c     distribute Zres(isle) to all perimeter points
c
         do isle=1,nisle
           if (imask(isle)) then
             do n=1,nippts(isle)
               i = iperm(iofs(isle)+n)
               j = jperm(iofs(isle)+n)
               Zres(i,j) = Zresisle(isle)
             end do
           end if
         end do
      end subroutine sum_dist


      subroutine avg_dist (Zres)
c
c     avg contributions to Zres(isle)
c     distribute Zres(isle) to all perimeter points
c
c     this subroutine converts a type(res_type) vector with
c     distributed contributions to perimeter values
c        Zres(isle) = avg (Zres(i,j))
c     into a type (dpsi_type) vector with replicated values
c     for land_mass perimeters
c        Zres(isle) = Zres(i,j)
c     for all (i,j) in the ocean perimeter of land_mass(isle).
c
c
      use congrad_module
      implicit none
      real(kind=8) Zres(imt,jmt), Zresisle(mnisle)
      integer i,j,isle,n
c
c     avg contributions to Zres(isle)
c
      do isle=1,nisle
        if (imask(isle)) then
          Zresisle(isle) = 0.0
          do n=1,nippts(isle)
            i = iperm(iofs(isle)+n)
            j = jperm(iofs(isle)+n)
            Zresisle(isle) = Zresisle(isle) + Zres(i,j)
c            if (isle>1) print*,isle,i,j,Zres(i,j),Zresisle(isle)
          end do
        end if
      end do
c
c     distribute Zres(isle) to all perimeter points
c
      do isle=1,nisle
        if (imask(isle)) then
          Zresisle(isle) = Zresisle(isle)/nippts(isle)
c          print*,isle,nippts(isle),Zresisle(isle)
          do n=1,nippts(isle)
            i = iperm(iofs(isle)+n)
            j = jperm(iofs(isle)+n)
            Zres(i,j) = Zresisle(isle)
          end do
        end if
      end do
      end subroutine avg_dist








