SUBMODULE(compute_caldyn_mod) compute_caldyn_vert_mod
   USE prec, ONLY: rstd
   USE ISO_C_BINDING, only: C_DOUBLE, C_FLOAT

   IMPLICIT NONE
   SAVE

CONTAINS

   MODULE SUBROUTINE compute_caldyn_vert(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du)
      USE grid_param, ONLY: grid_type, grid_ico
      USE disvert_mod, ONLY: bp
      USE trace
      REAL(rstd), INTENT(IN)    :: u(:, :)              ! momentum (covariant)
      REAL(rstd), INTENT(IN)    :: theta(:, :, :)        ! dynamical tracers (scalars)
      REAL(rstd), INTENT(IN)    :: rhodz(:, :)          ! mass (density)
      REAL(rstd), INTENT(INOUT) :: convm(:, :)          ! horizontal mass flux convergence (density)
      REAL(rstd), INTENT(INOUT) :: wflux(:, :)          ! vertical mass flux
      REAL(rstd), INTENT(OUT)   :: wwuu(:, :)           ! wflux * u
      REAL(rstd), INTENT(INOUT)  :: du(:, :)             ! momentum tendency
      REAL(rstd), INTENT(INOUT)  :: dtheta_rhodz(:, :, :) ! dynamical tracers tendencies (densities)
      REAL(rstd), INTENT(OUT)    :: dps(:)              ! column mass tendency

      CALL trace_start("compute_caldyn_vert")
      IF (grid_type == grid_ico) THEN
         CALL compute_caldyn_vert_hex_(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du, bp)
      ELSE
         CALL compute_caldyn_vert_unst_(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du)
      END IF
      CALL trace_end("compute_caldyn_vert")

   END SUBROUTINE compute_caldyn_vert

   !--------------------------------------------------------------

   SUBROUTINE compute_caldyn_vert_hex_(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du, bp_)
      USE grid_param
      USE disvert_mod
      USE caldyn_vars_mod
      USE omp_para
      REAL(rstd), INTENT(IN)  :: u(iim*3*jjm, llm)
      REAL(rstd), INTENT(IN)  :: theta(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(IN)  :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(INOUT)  :: convm(iim*jjm, llm)  ! mass flux convergence
      REAL(rstd), INTENT(INOUT) :: wflux(iim*jjm, llm + 1) ! vertical mass flux (kg/m2/s)
      REAL(rstd), INTENT(OUT)   :: wwuu(iim*3*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: du(iim*3*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: dtheta_rhodz(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(OUT) :: dps(iim*jjm)
      REAL(rstd), INTENT(IN)  :: bp_(llm)

      ! temporary variable
      INTEGER :: i, j, ij, l, iq
      REAL(rstd) :: p_ik, exner_ik, dF_deta, dFu_deta
      INTEGER    :: ij_omp_begin, ij_omp_end

      !$acc data async &
      !$acc present(rhodz(:,:), u(:,:), wwuu(:,:), wflux(:,:), dps(:), convm(:,:), du(:,:), dtheta_rhodz(:,:,:), theta(:,:,:), bp_(:))

!$OMP BARRIER

      CALL distrib_level(ij_begin, ij_end, ij_omp_begin, ij_omp_end)

      ! We use kernels here instead of "loop seq" + "loop gang vector"
      ! to be sure the code really abides the standard. In practice,
      ! it seems like the compiler interchanges the loops.

      !$acc kernels async
      DO l = llm - 1, 1, -1
         !DIR$ SIMD
         DO ij = ij_omp_begin, ij_omp_end
            ! cumulate mass flux convergence from top to bottom
            convm(ij, l) = convm(ij, l) + convm(ij, l + 1)
         END DO
      END DO
      !$acc end kernels

      !$OMP BARRIER
      ! FLUSH on convm

      IF (is_omp_first_level) THEN
         !$acc parallel loop  async
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dps(ij) = convm(ij, 1)
         END DO
      END IF

      !!! Compute vertical mass flux (l=1,llm+1 done by caldyn_BC)
      !$acc parallel loop collapse(2) async
      DO l = ll_beginp1, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            ! w = int(z,ztop,div(flux)dz) + B(eta)dps/dt
            ! => w>0 for upward transport
            wflux(ij, l) = bp_(l)*convm(ij, 1) - convm(ij, l)
         END DO
      END DO

      !--> flush wflux
      !$OMP BARRIER
      !$acc parallel loop collapse(3) async
      DO iq = 1, nqdyn
         DO l = ll_begin, ll_endm1
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq) - 0.5* &
                                         (wflux(ij, l + 1)*(theta(ij, l, iq) + theta(ij, l + 1, iq)))
            END DO
         END DO
      END DO

      !$acc parallel loop collapse(3) async
      DO iq = 1, nqdyn
         DO l = ll_beginp1, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq) + 0.5* &
                                         (wflux(ij, l)*(theta(ij, l - 1, iq) + theta(ij, l, iq)))
            END DO
         END DO
      END DO
      ! Compute vertical transport
      !$acc parallel loop collapse(2) async
      DO l = ll_beginp1, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            wwuu(ij + u_right, l) = 0.5*(wflux(ij, l) + wflux(ij + t_right, l))*(u(ij + u_right, l) - u(ij + u_right, l - 1))
            wwuu(ij + u_lup, l) = 0.5*(wflux(ij, l) + wflux(ij + t_lup, l))*(u(ij + u_lup, l) - u(ij + u_lup, l - 1))
            wwuu(ij + u_ldown, l) = 0.5*(wflux(ij, l) + wflux(ij + t_ldown, l))*(u(ij + u_ldown, l) - u(ij + u_ldown, l - 1))
         END DO
      END DO

      !--> flush wwuu
      !$OMP BARRIER

      ! Add vertical transport to du
      !$acc parallel loop collapse(2) async
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            du(ij + u_right, l) = du(ij + u_right, l) - &
            &                     (wwuu(ij + u_right, l + 1) + wwuu(ij + u_right, l))/(rhodz(ij, l) + rhodz(ij + t_right, l))
            du(ij + u_lup, l) = du(ij + u_lup, l) - &
            &                   (wwuu(ij + u_lup, l + 1) + wwuu(ij + u_lup, l))/(rhodz(ij, l) + rhodz(ij + t_lup, l))
            du(ij + u_ldown, l) = du(ij + u_ldown, l) - &
            &                     (wwuu(ij + u_ldown, l + 1) + wwuu(ij + u_ldown, l))/(rhodz(ij, l) + rhodz(ij + t_ldown, l))
         END DO
      END DO

      !$acc end data

   END SUBROUTINE compute_caldyn_vert_hex_

   SUBROUTINE compute_caldyn_vert_unst_(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du)
      USE data_unstructured_mod, ONLY: left, right, primal_num, edge_num
      USE disvert_mod, ONLY: llm, bp
      USE grid_param, ONLY: nqdyn
      USE caldyn_vars_mod
      USE omp_para
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)     :: u, du                 ! IN, INOUT
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn) :: theta, dtheta_rhodz   ! IN, INOUT
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)  :: rhodz, convm          ! IN, INOUT
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: wflux     ! BUF
      REAL(C_DOUBLE), DIMENSION(llm + 1, edge_num)    :: wwuu                  ! BUF
      REAL(C_DOUBLE), DIMENSION(primal_num)    :: dps                   ! OUT
      INTEGER :: ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown
      REAL(C_DOUBLE) :: sign1, sign2, sign3, sign4, sign5, sign6, le_de1, le_de2, le_de3, le_de4, le_de5, le_de6, &
                        de1, de2, de3, de4, de5, de6, le1, le2, le3, le4, le5, le6
      INTEGER :: edge1, edge2, edge3, edge4, edge5, edge6, ij_up1, ij_up2, ij_up3, ij_up4, ij_up5, ij_up6, &
                 ij_down1, ij_down2, ij_down3, ij_down4, ij_down5, ij_down6
      REAL(C_DOUBLE) :: dF_deta, dFu_deta

      wflux(:, :) = 0 ! FIXME, this should not be needed

      !--------------------------------------------------------------------------
      !---------------------------- caldyn_wflux ----------------------------------
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         DO l = llm - 1, 1, -1
            ! cumulate mass flux convergence from top to bottom
            convm(l, ij) = convm(l, ij) + convm(l + 1, ij)
         END DO
         dps(ij) = convm(1, ij)
         DO l = 2, llm
            ! Compute vertical mass flux (l=1,llm+1 set to zero at init)
            wflux(l, ij) = bp(l)*dps(ij) - convm(l, ij)
         END DO
      END DO
      !$OMP END DO
      ! make sure wflux is up to date
      !$OMP BARRIER
      !---------------------------- caldyn_wflux ----------------------------------
      !--------------------------------------------------------------------------

      !--------------------------------------------------------------------------
      !---------------------------- caldyn_vert ----------------------------------
      DO iq = 1, nqdyn
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 2, llm
               dtheta_rhodz(l, ij, iq) = dtheta_rhodz(l, ij, iq) + 0.5*(theta(l, ij, iq) + theta(l - 1, ij, iq))*wflux(l, ij)
            END DO
         END DO
         !$OMP END DO
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm - 1
               dtheta_rhodz(l, ij, iq) = dtheta_rhodz(l, ij, iq) - 0.5*(theta(l, ij, iq) + theta(l + 1, ij, iq))*wflux(l + 1, ij)
            END DO
         END DO
         !$OMP END DO
      END DO
      IF (caldyn_vert_variant == caldyn_vert_cons) THEN
         ! conservative vertical transport of momentum : (F/m)du/deta = 1/m (d/deta(Fu)-u.dF/deta)
         !$OMP DO SCHEDULE(STATIC)
         DO edge = 1, edge_num
            ij_left = left(edge)
            ij_right = right(edge)
            !DIR$ SIMD
            DO l = 2, llm
               wwuu(l, edge) = .25*(wflux(l, ij_left) + wflux(l, ij_right))*(u(l, edge) + u(l - 1, edge)) ! Fu
            END DO
         END DO
         !$OMP END DO
         ! make sure wwuu is up to date
         !$OMP BARRIER
         !$OMP DO SCHEDULE(STATIC)
         DO edge = 1, edge_num
            ij_left = left(edge)
            ij_right = right(edge)
            !DIR$ SIMD
            DO l = 1, llm
               dFu_deta = wwuu(l + 1, edge) - wwuu(l, edge) ! d/deta (F*u)
               dF_deta = .5*(wflux(l + 1, ij_left) + wflux(l + 1, ij_right) - (wflux(l, ij_left) + wflux(l, ij_right))) ! d/deta(F)
               du(l, edge) = du(l, edge) - (dFu_deta - u(l, edge)*dF_deta)/(.5*(rhodz(l, ij_left) + rhodz(l, ij_right))) ! (F/m)du/deta
            END DO
         END DO
         !$OMP END DO
      ELSE
         !$OMP DO SCHEDULE(STATIC)
         DO edge = 1, edge_num
            ij_left = left(edge)
            ij_right = right(edge)
            !DIR$ SIMD
            DO l = 2, llm
               wwuu(l, edge) = .5*(wflux(l, ij_left) + wflux(l, ij_right))*(u(l, edge) - u(l - 1, edge))
            END DO
         END DO
         !$OMP END DO
         ! make sure wwuu is up to date
         !$OMP BARRIER
         !$OMP DO SCHEDULE(STATIC)
         DO edge = 1, edge_num
            ij_left = left(edge)
            ij_right = right(edge)
            !DIR$ SIMD
            DO l = 1, llm
               du(l, edge) = du(l, edge) - (wwuu(l, edge) + wwuu(l + 1, edge))/(rhodz(l, ij_left) + rhodz(l, ij_right))
            END DO
         END DO
         !$OMP END DO
      END IF
      !---------------------------- caldyn_vert ----------------------------------
      !--------------------------------------------------------------------------

   END SUBROUTINE compute_caldyn_vert_unst_

   SUBROUTINE compute_caldyn_vert_manual(u, theta, rhodz, convm, wflux, wwuu, dps, dtheta_rhodz, du)
      USE grid_param
      USE disvert_mod
      USE caldyn_vars_mod
      USE omp_para
      REAL(rstd), INTENT(IN)  :: u(iim*3*jjm, llm)
      REAL(rstd), INTENT(IN)  :: theta(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(IN)  :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: convm(iim*jjm, llm)  ! mass flux convergence
      REAL(rstd), INTENT(INOUT) :: wflux(iim*jjm, llm + 1) ! vertical mass flux (kg/m2/s)
      REAL(rstd), INTENT(OUT)   :: wwuu(iim*3*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: du(iim*3*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: dtheta_rhodz(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(OUT) :: dps(iim*jjm)

      ! temporary variable
      INTEGER :: i, j, ij, l, iq
      REAL(rstd) :: p_ik, exner_ik, dF_deta, dFu_deta
      INTEGER    :: ij_omp_begin, ij_omp_end

!$OMP BARRIER

      CALL distrib_level(ij_begin, ij_end, ij_omp_begin, ij_omp_end)

!!! cumulate mass flux convergence from top to bottom
      DO l = llm - 1, 1, -1
         !DIR$ SIMD
         DO ij = ij_omp_begin, ij_omp_end
            convm(ij, l) = convm(ij, l) + convm(ij, l + 1)
         END DO
      END DO
      !  ENDIF

      !$OMP BARRIER
      ! FLUSH on convm
      ! compute dmass_col
      IF (is_omp_first_level) THEN
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            ! dps/dt = -int(div flux)dz
            dps(ij) = convm(ij, 1)
         END DO
      END IF
!!! Compute vertical mass flux (l=1,llm+1 done by caldyn_BC)
      DO l = ll_beginp1, ll_end
         !    IF (caldyn_conserv==energy) CALL test_message(req_qu)
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            ! w = int(z,ztop,div(flux)dz) + B(eta)dps/dt
            ! => w>0 for upward transport
            wflux(ij, l) = bp(l)*convm(ij, 1) - convm(ij, l)
         END DO
      END DO

      !--> flush wflux
      !$OMP BARRIER

      DO iq = 1, nqdyn
         DO l = ll_begin, ll_endm1
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq) - 0.5* &
                                         (wflux(ij, l + 1)*(theta(ij, l, iq) + theta(ij, l + 1, iq)))
            END DO
         END DO
         DO l = ll_beginp1, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq) + 0.5* &
                                         (wflux(ij, l)*(theta(ij, l - 1, iq) + theta(ij, l, iq)))
            END DO

         END DO
      END DO

      ! Compute vertical transport
      DO l = ll_beginp1, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            wwuu(ij + u_right, l) = 0.5*(wflux(ij, l) + wflux(ij + t_right, l))*(u(ij + u_right, l) - u(ij + u_right, l - 1))
            wwuu(ij + u_lup, l) = 0.5*(wflux(ij, l) + wflux(ij + t_lup, l))*(u(ij + u_lup, l) - u(ij + u_lup, l - 1))
            wwuu(ij + u_ldown, l) = 0.5*(wflux(ij, l) + wflux(ij + t_ldown, l))*(u(ij + u_ldown, l) - u(ij + u_ldown, l - 1))
         END DO
      END DO

      !--> flush wwuu
      !$OMP BARRIER

      ! Add vertical transport to du
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            du(ij + u_right, l) = du(ij + u_right, l) - &
                                  (wwuu(ij + u_right, l + 1) + wwuu(ij + u_right, l))/(rhodz(ij, l) + rhodz(ij + t_right, l))
            du(ij + u_lup, l) = du(ij + u_lup, l) - &
                                (wwuu(ij + u_lup, l + 1) + wwuu(ij + u_lup, l))/(rhodz(ij, l) + rhodz(ij + t_lup, l))
            du(ij + u_ldown, l) = du(ij + u_ldown, l) - &
                                  (wwuu(ij + u_ldown, l + 1) + wwuu(ij + u_ldown, l))/(rhodz(ij, l) + rhodz(ij + t_ldown, l))
         END DO
      END DO

   END SUBROUTINE compute_caldyn_vert_manual

END SUBMODULE compute_caldyn_vert_mod
