SUBMODULE(compute_transport_mod) compute_vlz_mod

   IMPLICIT NONE

   REAL(rstd), PARAMETER :: pente_max = 2.0 ! for vlz

CONTAINS

   MODULE SUBROUTINE compute_vlz(update_mass, fac, halo, massp, fluxp, massq, fluxq, q, dzqw, adzqw, dzq)
      USE grid_param, ONLY: grid_type, grid_ico
      LOGICAL, INTENT(IN)       :: update_mass
      REAL(rstd), INTENT(IN)    :: fac
      INTEGER, INTENT(IN)       :: halo
      REAL(rstd), INTENT(INOUT) :: massp(:, :), massq(:, :) ! mass of parent (p) and tracer (q)
      REAL(rstd), INTENT(INOUT) :: fluxp(:, :), fluxq(:, :)  ! time-integrated flux of parent (p) and tracer (q)
      REAL(rstd), INTENT(INOUT) :: q(:, :)                 ! mass ratio of tracer
      ! temporary fields, need to be shared among OpenMP threads
      REAL(rstd), INTENT(INOUT) :: dzqw(:, :), &  ! vertical finite difference of q
      &                            adzqw(:, :), & ! abs(dzqw)
      &                            dzq(:, :)      ! limited slope of q

      IF (grid_type == grid_ico) THEN
         CALL compute_vlz_hex(update_mass, fac, halo, &
                              massp(:, :), fluxp(:, :), &
                              massq(:, :), fluxq(:, :), &
                              q, dzqw, adzqw, dzq)
      ELSE
         CALL compute_vlz_unst(update_mass, fac, &
                               massp(:, :), fluxp(:, :), &
                               massq(:, :), fluxq(:, :), &
                               q, dzqw, dzq)
      END IF
   END SUBROUTINE

   SUBROUTINE compute_vlz_hex(update_mass, fac, halo, massp, fluxp, massq, fluxq, q, dzqw, adzqw, dzq)
      !
      !     Auteurs:   P.Le Van, F.Hourdin, F.Forget, T. Dubos
      !
      !    ********************************************************************
      !     Update tracers using vertical mass flux only
      !     Van Leer scheme with minmod limiter
      !     fluxp >0 for upward transport
      !    ********************************************************************
      USE trace, ONLY: trace_start, trace_end
      USE omp_para, ONLY: ll_begin, ll_end, ll_beginp1, ll_endp1, ll_endm1, &
      &                   is_omp_first_level, is_omp_last_level
      USE icosa, ONLY: nqtot, llm, iim, jjm, ii_begin, ii_end, jj_begin, jj_end

      LOGICAL, INTENT(IN)       :: update_mass
      REAL(rstd), INTENT(IN)    :: fac
      INTEGER, INTENT(IN)       :: halo
      REAL(rstd), INTENT(INOUT) :: fluxp(iim*jjm, llm + 1), fluxq(iim*jjm, llm + 1)  ! time-integrated fluxes
      REAL(rstd), INTENT(INOUT) :: massp(iim*jjm, llm), massq(iim*jjm, llm)  ! mass
      REAL(rstd), INTENT(INOUT) :: q(iim*jjm, llm)

      ! temporary shared variable
      REAL(rstd), INTENT(INOUT) :: dzqw(iim*jjm, llm), & ! vertical finite difference of q
                                   adzqw(iim*jjm, llm), & ! abs(dzqw)
                                   dzq(iim*jjm, llm)           ! limited slope of q

      REAL(rstd) :: dzqmax, newmass, sigw, qq, w
      INTEGER :: ij, l, ijb, ije

      CALL trace_start("vlz")

      ijb = ((jj_begin - halo) - 1)*iim + ii_begin - halo
      ije = ((jj_end + halo) - 1)*iim + ii_end + halo

      ! finite difference of q

      !$acc data present(q(:,:), massp(:,:), massq(:,:), dzqw(:,:), adzqw(:,:), dzq(:,:), fluxp(:,:), fluxq(:,:)) async

      !$acc parallel loop async collapse(2)
      DO l = ll_beginp1, ll_end
         !DIR$ SIMD
         DO ij = ijb, ije
            dzqw(ij, l) = q(ij, l) - q(ij, l - 1)
            adzqw(ij, l) = abs(dzqw(ij, l))
         END DO
      END DO

      !--> flush dzqw, adzqw
      !$OMP BARRIER

      ! minmod-limited slope of q
      ! dzq = slope*dz, i.e. the reconstructed q varies by dzq inside level l

      !$acc parallel loop async collapse(2)
      DO l = ll_beginp1, ll_endm1
         !DIR$ SIMD
         DO ij = ijb, ije
            IF (dzqw(ij, l)*dzqw(ij, l + 1) .gt. 0.) THEN
               dzq(ij, l) = 0.5*(dzqw(ij, l) + dzqw(ij, l + 1))
               dzqmax = pente_max*min(adzqw(ij, l), adzqw(ij, l + 1))
               dzq(ij, l) = sign(min(abs(dzq(ij, l)), dzqmax), dzq(ij, l))  ! NB : sign(a,b)=a*sign(b)
            ELSE
               dzq(ij, l) = 0.
            END IF
         END DO
      END DO

      ! 0 slope in top and bottom layers
      IF (is_omp_first_level) THEN
         !$acc parallel loop async
         DO ij = ijb, ije
            dzq(ij, 1) = 0.
         END DO
      END IF

      IF (is_omp_last_level) THEN
         !$acc parallel loop async
         DO ij = ijb, ije
            dzq(ij, llm) = 0.
         END DO
      END IF

      !---> flush dzq
      !$OMP BARRIER

      ! sigw = fraction of mass that leaves level l/l+1
      ! then amount of q leaving level l/l+1 = fluxq = fluxp * qq
      !$acc parallel loop collapse(2) async
      DO l = ll_beginp1, ll_end
         !DIR$ SIMD
         DO ij = ijb, ije
            w = fac*fluxp(ij, l)
            IF (w > 0.) THEN  ! upward transport, upwind side is at level l
               sigw = w/massp(ij, l - 1)
               qq = q(ij, l - 1) + 0.5*(1.-sigw)*dzq(ij, l - 1) ! qq = q if sigw=1 , qq = q+dzq/2 if sigw=0
            ELSE           ! downward transport, upwind side is at level l+1
               sigw = w/massp(ij, l)
               qq = q(ij, l) - 0.5*(1.+sigw)*dzq(ij, l) ! qq = q if sigw=-1 , qq = q-dzq/2 if sigw=0
            END IF
            fluxq(ij, l) = w*qq
         END DO
      END DO
      ! fluxq = 0 at top and bottom
      IF (is_omp_first_level) THEN
         !$acc parallel loop async
         DO ij = ijb, ije
            fluxq(ij, 1) = 0.
         END DO
      END IF

      IF (is_omp_last_level) THEN
         !$acc parallel loop async
         DO ij = ijb, ije
            fluxq(ij, llm + 1) = 0.
         END DO
      END IF

      ! --> flush fluxq
      !$OMP BARRIER

      ! update q, mass is updated only after all q's have been updated
      !$acc parallel loop collapse(2) async
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ijb, ije
            massq(ij, l) = q(ij, l)*massp(ij, l)
            newmass = massp(ij, l) + fac*(fluxp(ij, l) - fluxp(ij, l + 1))
            q(ij, l) = (massq(ij, l) + fluxq(ij, l) - fluxq(ij, l + 1))/newmass
            IF (update_mass) massp(ij, l) = newmass
         END DO
      END DO
      !$acc end data
      CALL trace_end("vlz")

   END SUBROUTINE compute_vlz_hex

   !====================================== Kernels ========================================

   SUBROUTINE compute_vlz_unst(update_mass, fac, massp, fluxp, massq, fluxq, q, dzqw, dzq)
      USE trace, ONLY: trace_start, trace_end
      USE grid_param, ONLY: llm, primal_num
      LOGICAL, INTENT(IN)    :: update_mass
      REAL(rstd), INTENT(IN)    :: fac
      REAL(rstd), DIMENSION(llm + 1, primal_num), INTENT(IN) :: fluxp
      REAL(rstd), DIMENSION(llm + 1, primal_num), INTENT(OUT) :: fluxq
      REAL(rstd), DIMENSION(llm, primal_num), INTENT(INOUT)  :: massp, massq, q
      REAL(rstd), DIMENSION(llm, primal_num), INTENT(OUT)    :: dzqw, dzq

      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown
      REAL(rstd) :: dzqm, dzqmax, sigw, w, qq, newmass

      CALL trace_start("vlz")

      !--------------------------------------------------------------------------
      !---------------------------- vlz ----------------------------------
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         !DIR$ SIMD
         DO l = 2, llm
            dzqw(l, ij) = q(l, ij) - q(l - 1, ij)
         END DO
      END DO
      !$OMP END DO
      !
      !$OMP BARRIER
      !
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         l = 1
         dzq(l, ij) = 0. ! zero slope in bottom layer
         !DIR$ SIMD
         DO l = 2, llm - 1
            IF (dzqw(l + 1, ij)*dzqw(l, ij) > 0.) THEN
               dzqm = .5*(dzqw(l + 1, ij) + dzqw(l, ij))
               dzqmax = pente_max*MIN(ABS(dzqw(l + 1, ij)), ABS(dzqw(l, ij)))
               ! NB : sign(a,b) = a*sign(b)
               dzq(l, ij) = sign(MIN(ABS(dzqm), dzqmax), dzqm)
            ELSE
               dzq(l, ij) = 0.
            END IF
         END DO
         l = llm
         dzq(l, ij) = 0. ! zero slope in top layer
      END DO
      !$OMP END DO
      !
      !$OMP BARRIER
      !
      ! w : fraction of massp that leaves level l-1 (fluxp>0) or l (fluxp<0)
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         l = 1
         fluxq(l, ij) = 0.
         !DIR$ SIMD
         DO l = 2, llm
            w = fac*fluxp(l, ij)
            IF (w > 0.) THEN
               ! upward transport, upwind side is at level l
               ! qq = q if sigw=1 , qq = q+dzq/2 if sigw=0
               sigw = w/massp(l - 1, ij)
               qq = q(l - 1, ij) + 0.5*(1.-sigw)*dzq(l - 1, ij)
            ELSE
               ! downward transport, upwind side is at level l+1
               ! qq = q if sigw=-1 , qq = q-dzq/2 if sigw=0
               sigw = w/massp(l, ij)
               qq = q(l, ij) - 0.5*(1.+sigw)*dzq(l, ij)
            END IF
            fluxq(l, ij) = w*qq
         END DO
         l = llm + 1
         fluxq(l, ij) = 0.
      END DO
      !$OMP END DO
      !
      !$OMP BARRIER
      !
      ! update q, massp is updated only after all q s have been updated
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         !DIR$ SIMD
         DO l = 1, llm
            massq(l, ij) = q(l, ij)*massp(l, ij)
            newmass = massp(l, ij) + fac*(fluxp(l, ij) - fluxp(l + 1, ij))
            q(l, ij) = (massq(l, ij) + fluxq(l, ij) - fluxq(l + 1, ij))/newmass
            IF (update_mass) massp(l, ij) = newmass
         END DO
      END DO
      !$OMP END DO
      !---------------------------- vlz ----------------------------------
      !--------------------------------------------------------------------------
      CALL trace_end("vlz")

   END SUBROUTINE compute_vlz_unst

END SUBMODULE compute_vlz_mod
