#include "options.inc"


c
c-----------------------------------------------------------------------
c     implicit vertical mixing module
c
c     linked in the code in setup, tracer and clinic
c
c                     c eden
c-----------------------------------------------------------------------
c
      subroutine implicitvmix_init
      use spflame_module
      implicit none

      if (my_pe==0) then
       print*,' Initialisation of implicit vertical mixing'
      endif
     
      if (enable_implicit_vert_diff) then
c         aidif_tracer = 0.5  ! Crank Nicholsen, semi implicit
         aidif_tracer = 1.0  
      else
         aidif_tracer = 0.0  ! explicit
      endif

      if (enable_implicit_vert_fric) then
c         aidif_momentum = 0.5
         aidif_momentum = 1.0
      else
         aidif_momentum = 0.0
      endif

      if (my_pe==0) then
       print*,' aidif_tracer  =',aidif_tracer
       print*,' aidif_momentum=',aidif_momentum
       print*,' 0.5 means semi-implicit, 0.0 explicit '
       print*,' done '
      endif

      end subroutine implicitvmix_init


      subroutine implicit_vert_diff(n,is,ie,js,je)
c
c-----------------------------------------------------------------------
c     solve vertical diffusion of tracers implicitly
c     author:   r.c.pacanowski       e-mail rcp@gfdl.gov
c     Klaus Ketelsen     29.1.99   Loop order changed
c     Klaus Ketelsen     02.2.99   setup of equation, and 1. step invert mixed
c-----------------------------------------------------------------------
c
      use spflame_module
      implicit none
      integer , intent(in) :: n,is,ie,js,je
#ifndef vector_host
c
c-----------------------------------------------------------------------
c     risc version of the code
c-----------------------------------------------------------------------
c
      integer              :: i,j,k,kz
      real                                     :: a,b,c,d,g,fxa
      real, dimension(0:km)                    :: e,f
# ifdef partial_cell
      real ::  dhw_ut(0:km)
      real ::  dhut  (  km)
# endif

      fxa = c2dt*aidif_tracer
c
c     add in the implicit vertical diffusion
c
      do j=js,je
       do i=is,ie
# ifdef partial_cell
          do k=1,km
           dhut(k) = dht(i,k,j)
	  enddo
          do k=0,km
	   dhw_ut(k) = dhwt(i,k,j)
	  enddo
# endif
        kz = kmt(i,j)
        do k=kz+1,km
           e(k-1) = 0.
           f(k-1) = 0.
        end do
!       b. c. at bottom
        if (kz .ne. 0) then
# ifdef partial_cell
          a     = diff_cbt(i,kz-1,j)*fxa/(dhw_ut(kz-1)*dhut(kz))
          c     = fxa/dhut(kz)
# else
          a     = diff_cbt(i,kz-1,j)*dztur(kz)*fxa
          c     = dztr(kz)*fxa
# endif
          b     = 1. + a
          e     = 0.
          d     = t(i,kz,j,n,taup1)
          f(kz) = -btf(i,j,n)
          k = kz 
          g      = 1./(b-c*e(k))
          e(k-1) = a*g
          f(k-1) = (d+c*f(k))*g
        endif
!       interior
        do k=kz-1,2,-1
# ifdef partial_cell
           a      = diff_cbt(i,k-1,j)*fxa/(dhw_ut(k-1)*dhut(k))
           c      = diff_cbt(i,k,j)*fxa/(dhw_ut(k)*dhut(k)) 
# else
           a      = diff_cbt(i,k-1,j)*dztur(k)*fxa
           c      = diff_cbt(i,k,j)*dztlr(k)*fxa
# endif
           b      = 1. + a + c
           d      = t(i,k,j,n,taup1)
           g      = 1./(b-c*e(k))
           e(k-1) = a*g
           f(k-1) = (d+c*f(k))*g
        enddo
!       b. c. at top
        if(kz.ne.1)   then
           k = 1
# ifdef partial_cell
           a      = fxa/dhut(k)
           c      = diff_cbt(i,k,j)*fxa/(dhw_ut(k)*dhut(k))
# else
           a      = dztr(k)*fxa
           c      = diff_cbt(i,k,j)*dztlr(k)*fxa
# endif
           b      = 1. + c
           d      = t(i,k,j,n,taup1)
!       now invert
           g      = 1./(b-c*e(k))
           e(k-1) = a*g
           f(k-1) = (d+c*f(k))*g
        end if
!       b.c. at surface
        t(i,1,j,n,taup1) = (e(0)*stf(i,j,n) + f(0))*tmask(i,1,j)
        do k=2,km
         t(i,k,j,n,taup1) = (e(k-1)*t(i,k-1,j,n,taup1) 
     &                       + f(k-1))*tmask(i,k,j)
        enddo
       enddo    ! end i-loop
      enddo     ! end j-loop
#else
c
c-----------------------------------------------------------------------
c     vector version of the code
c-----------------------------------------------------------------------
c
      real, dimension(is:ie,km)   :: a,b,c,d
      real, dimension(is:ie,0:km) :: e,f
      real, dimension(is:ie)      :: g
      integer i,j,k,kz
      real fxa
c
      fxa = c2dt*aidif_tracer
c
      do j=js,je
       do k=2,km
          do i=is,ie
#ifdef partial_cell
            a(i,k)   = diff_cbt(i,k-1,j)*fxa/(dhwt(i,k-1,j)*dht(i,k,j))
            c(i,k)   = diff_cbt(i,k,j)*fxa/(dhwt(i,k,j)*dht(i,k,j))
#else
            a(i,k)   = diff_cbt(i,k-1,j)*dztur(k)*fxa
            c(i,k)   = diff_cbt(i,k,j)*dztlr(k)*fxa
#endif
            b(i,k)   = 1. + a(i,k) + c(i,k)
            d(i,k)   = t(i,k,j,n,taup1)
            e(i,k-1) = 0.
            f(i,k-1) = 0.
          enddo
       enddo
c      b. c. at top
       k = 1
       do i=is,ie
#ifdef partial_cell
          a(i,k)   = fxa/dht(i,k,j)
          c(i,k)   = diff_cbt(i,k,j)*fxa/(dhwt(i,k,j)*dht(i,k,j))
#else
          a(i,k)   = dztr(k)*fxa
          c(i,k)   = diff_cbt(i,k,j)*dztlr(k)*fxa
#endif
          b(i,k)   = 1. + c(i,k)
          d(i,k)   = t(i,k,j,n,taup1)
          e(i,k-1) = 0.
          f(i,k-1) = 0.
       enddo
c      b. c. at bottom
       do i=is,ie
         kz = kmt(i,j)
         if (kz .ne. 0) then
            b(i,kz) = 1. + a(i,kz)
#ifdef partial_cell
            c(i,kz) = fxa/dht(i,kz,j)
#else
            c(i,kz) = dztr(kz)*fxa
#endif
            e(i,kz) = 0.
            f(i,kz) = -btf(i,j,n)
         endif
       enddo
c      now invert
       do k=km,1,-1
         do i=is,ie
          if (k .le. kmt(i,j)) then
              g(i)     = 1./(b(i,k)-c(i,k)*e(i,k))
              e(i,k-1) = a(i,k)*g(i)
              f(i,k-1) = (d(i,k)+c(i,k)*f(i,k))*g(i)
          endif
         enddo
       enddo
c      b.c. at surface
       do i=is,ie
         t(i,1,j,n,taup1) = (e(i,0)*stf(i,j,n) + f(i,0))*tmask(i,1,j)
       enddo
       do k=2,km
        do i=is,ie
          t(i,k,j,n,taup1) = (e(i,k-1)*t(i,k-1,j,n,taup1) 
     &                        + f(i,k-1))*tmask(i,k,j)
        enddo
       enddo
      enddo ! j-loop
#endif
      end subroutine implicit_vert_diff



      subroutine implicit_vert_fric(n,is,ie,js,je)
c
c-----------------------------------------------------------------------
c     solve vertical diffusion of velocity implicitly
c     author:   r.c.pacanowski       e-mail rcp@gfdl.gov
c     Klaus Ketelsen     29.1.99   Loop order changed
c     Klaus Ketelsen     02.2.99   setup of equation, and 1. step invert mixed
c-----------------------------------------------------------------------
c
      use spflame_module
      implicit none
      integer , intent(in) :: is,ie,js,je,n
      real                 :: fxa
      integer              :: i,j,k,kz
#ifndef vector_host
c
c-----------------------------------------------------------------------
c     risc version arrays
c-----------------------------------------------------------------------
c
      real                                     :: a,b,c,d
      real, dimension(0:km)                    :: e,f
      real                                     :: g
# ifdef partial_cell
      real ::  dhw_ut(0:km)
      real ::  dhut  (  km)
# endif
#else
c
c-----------------------------------------------------------------------
c     vector version arrays
c-----------------------------------------------------------------------
c
      real, dimension(is:ie,km)   :: a,b,c,d
      real, dimension(is:ie,0:km) :: e,f
      real, dimension(is:ie)      :: g
# ifdef partial_cell
      real ::  dhw_ut(is:ie,0:km)
# endif
#endif
c
c     construct the "tau+1" velocity without implicit vertical diffusion
c
      do j=js,je
        do k=1,km
          do i=is,ie
            u(i,k,j,n,taup1) = u(i,k,j,n,taum1)+c2dt*u(i,k,j,n,taup1)
          enddo
        enddo
      enddo
      fxa = rho0r*c2dt*aidif_momentum
c
c     add in the implicit vertical diffusion
c
#ifndef vector_host
c
c-----------------------------------------------------------------------
c     Risc version of code
c-----------------------------------------------------------------------
c
      do j=js,je
       do i=is,ie
# ifdef partial_cell
        do k=1,km
	      dhut(k) = dhu(i,k,j)
        enddo
        do k=0,km
	      dhw_ut(k) = min(dhwt(i,k,j),   dhwt(i+1,k,j)
     &,                       dhwt(i,k,j+1), dhwt(i+1,k,j+1))
	enddo
# endif
        kz = kmu(i,j)
        do k=kz+1,km
           e(k-1) = 0.
           f(k-1) = 0.
        end do
!   b. c. at bottom
        if (kz .ne. 0) then
# ifdef partial_cell
          a     = visc_cbu(i,kz-1,j)*fxa/(dhw_ut(kz-1)*dhut(kz))
          c     = fxa/dhut(kz)
# else
          a     = visc_cbu(i,kz-1,j)*dztur(kz)*fxa
          c     = dztr(kz)*fxa
# endif
          b     = 1. + a
          e     = 0.
          d     = u(i,kz,j,n,taup1)
          f(kz) = -bmf(i,j,n)
          k = kz 
          g      = 1./(b-c*e(k))
          e(k-1) = a*g
          f(k-1) = (d+c*f(k))*g
        endif
!       interior
        do k=kz-1,2,-1
# ifdef partial_cell
           a      = visc_cbu(i,k-1,j)*fxa/(dhw_ut(k-1)*dhut(k))
           c      = visc_cbu(i,k,j)*fxa/(dhw_ut(k)*dhut(k)) 
# else
           a      = visc_cbu(i,k-1,j)*dztur(k)*fxa
           c      = visc_cbu(i,k,j)*dztlr(k)*fxa
# endif
           b      = 1. + a + c
           d      = u(i,k,j,n,taup1)
           g      = 1./(b-c*e(k))
           e(k-1) = a*g
           f(k-1) = (d+c*f(k))*g
        enddo
!     b. c. at top
        if(kz.ne.1)   then
           k = 1
# ifdef partial_cell
           a      = fxa/dhut(k)
           c      = visc_cbu(i,k,j)*fxa/(dhw_ut(k)*dhut(k))
# else
           a      = dztr(k)*fxa
           c      = visc_cbu(i,k,j)*dztlr(k)*fxa
# endif
           b      = 1. + c
           d      = u(i,k,j,n,taup1)
!     now invert
           g      = 1./(b-c*e(k))
           e(k-1) = a*g
           f(k-1) = (d+c*f(k))*g
        end if
!     b.c. at surface
        u(i,1,j,n,taup1) = (e(0)*smf(i,j,n) + f(0))*umask(i,1,j)
        do k=2,km
         u(i,k,j,n,taup1) = (e(k-1)*u(i,k-1,j,n,taup1) 
     &                       + f(k-1))*umask(i,k,j)
        enddo
       enddo    ! end i-loop
      enddo     ! end j-loop
#else
c
c-----------------------------------------------------------------------
c    vector version of code
c-----------------------------------------------------------------------
c
      do j=js,je

# ifdef partial_cell
       do k=0,km
        do i=is,ie
	    dhw_ut(i,k) = min(dhwt(i,k,j),   dhwt(i+1,k,j),
     &                        dhwt(i,k,j+1), dhwt(i+1,k,j+1))
        enddo
       enddo
# endif

       do k=2,km
          do i=is,ie
#ifdef partial_cell
            a(i,k)   = visc_cbu(i,k-1,j)*fxa/(dhw_ut(i,k-1)*dhu(i,k,j))
            c(i,k)   = visc_cbu(i,k,j)*fxa/(dhw_ut(i,k)*dhu(i,k,j))
#else
            a(i,k)   = visc_cbu(i,k-1,j)*dztur(k)*fxa
            c(i,k)   = visc_cbu(i,k,j)*dztlr(k)*fxa
#endif
            b(i,k)   = 1. + a(i,k) + c(i,k)
            d(i,k)   = u(i,k,j,n,taup1)
            e(i,k-1) = 0.
            f(i,k-1) = 0.
          enddo
       enddo
c      b. c. at top
       k = 1
       do i=is,ie
#ifdef partial_cell
          a(i,k)   = fxa/dhu(i,k,j)
          c(i,k)   = visc_cbu(i,k,j)*fxa/(dhw_ut(i,k)*dhu(i,k,j))
#else
          a(i,k)   = dztr(k)*fxa
          c(i,k)   = visc_cbu(i,k,j)*dztlr(k)*fxa
#endif
          b(i,k)   = 1. + c(i,k)
          d(i,k)   = u(i,k,j,n,taup1)
          e(i,k-1) = 0.
          f(i,k-1) = 0.
       enddo
c      b. c. at bottom
       do i=is,ie
         kz = kmu(i,j)
         if (kz .ne. 0) then
            b(i,kz) = 1. + a(i,kz)
#ifdef partial_cell
            c(i,kz) = fxa/dhu(i,kz,j)
#else
            c(i,kz) = dztr(kz)*fxa
#endif
            e(i,kz) = 0.
            f(i,kz) = -bmf(i,j,n)
         endif
       enddo
c      now invert
       do k=km,1,-1
         do i=is,ie
          if (k .le. kmu(i,j)) then
              g(i)     = 1./(b(i,k)-c(i,k)*e(i,k))
              e(i,k-1) = a(i,k)*g(i)
              f(i,k-1) = (d(i,k)+c(i,k)*f(i,k))*g(i)
          endif
         enddo
       enddo
c      b.c. at surface
       do i=is,ie
         u(i,1,j,n,taup1) = (e(i,0)*smf(i,j,n) + f(i,0))*umask(i,1,j)
       enddo
       do k=2,km
        do i=is,ie
          u(i,k,j,n,taup1) = (e(i,k-1)*u(i,k-1,j,n,taup1) 
     &                        + f(i,k-1))*umask(i,k,j)
        enddo
       enddo
      enddo ! j-loop
#endif
c
c-----------------------------------------------------------------------
c     convert back to time change of velocity
c-----------------------------------------------------------------------
c
      do j=js,je
        do k=1,km
          do i=is,ie
            u(i,k,j,n,taup1) =(u(i,k,j,n,taup1)-u(i,k,j,n,taum1))/c2dt
          enddo
        enddo
      enddo
      end subroutine implicit_vert_fric



