SUBMODULE(compute_caldyn_mod) compute_caldyn_vert_NH_mod
   USE prec, ONLY: rstd
   IMPLICIT NONE

CONTAINS

   MODULE SUBROUTINE compute_caldyn_vert_NH(mass, geopot, W, wflux, W_etadot, du, dPhi, dW)
      REAL(rstd), INTENT(IN) :: mass(:, :)
      REAL(rstd), INTENT(IN) :: geopot(:, :)
      REAL(rstd), INTENT(IN) :: W(:, :)
      REAL(rstd), INTENT(IN) :: wflux(:, :)
      REAL(rstd), INTENT(INOUT) :: du(:, :)
      REAL(rstd), INTENT(INOUT) :: dPhi(:, :)
      REAL(rstd), INTENT(INOUT) :: dW(:, :)
      REAL(rstd) :: W_etadot(:, :) ! vertical flux of vertical momentum
      CALL compute_caldyn_vert_NH_manual(mass, geopot, W, wflux, W_etadot, du, dPhi, dW)
   END SUBROUTINE

   SUBROUTINE compute_caldyn_vert_NH_manual(mass, geopot, W, wflux, W_etadot, du, dPhi, dW)
      USE dimensions ! iim, jjm, offsets ...
      USE geometry ! signs
      USE grid_param, ONLY: llm, nqdyn
      USE trace, ONLY: trace_start, trace_end
      USE omp_para, ONLY: ll_begin, ll_beginp1, ll_end, ll_endp1, ll_endm1
      REAL(rstd), INTENT(IN) :: mass(iim*jjm, llm)
      REAL(rstd), INTENT(IN) :: geopot(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN) :: W(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN) :: wflux(iim*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: du(iim*3*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: dPhi(iim*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: dW(iim*jjm, llm + 1)
      REAL(rstd) :: W_etadot(iim*jjm, llm) ! vertical flux of vertical momentum
      ! local arrays
      REAL(rstd) :: eta_dot(iim*jjm, llm) ! eta_dot in full layers
      REAL(rstd) :: wcov(iim*jjm, llm) ! covariant vertical momentum in full layers
      ! indices and temporary values
      INTEGER    :: ij, l
      REAL(rstd) :: wflux_ij, w_ij

      CALL trace_start("compute_caldyn_vert_nh")

      DO l = ll_begin, ll_end
         ! compute the local arrays
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            wflux_ij = .5*(wflux(ij, l) + wflux(ij, l + 1))
            w_ij = .5*(W(ij, l) + W(ij, l + 1))/mass(ij, l)
            W_etadot(ij, l) = wflux_ij*w_ij
            eta_dot(ij, 1) = wflux_ij/mass(ij, l)
            wcov(ij, 1) = w_ij*(geopot(ij, l + 1) - geopot(ij, l))
         END DO
         ! add NH term to du
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            du(ij + u_right, l) = du(ij + u_right, l) &
                                  - .5*(wcov(ij + t_right, 1) + wcov(ij, 1)) &
                                  *ne_right*(eta_dot(ij + t_right, 1) - eta_dot(ij, 1))
            du(ij + u_lup, l) = du(ij + u_lup, l) &
                                - .5*(wcov(ij + t_lup, 1) + wcov(ij, 1)) &
                                *ne_lup*(eta_dot(ij + t_lup, 1) - eta_dot(ij, 1))
            du(ij + u_ldown, l) = du(ij + u_ldown, l) &
                                  - .5*(wcov(ij + t_ldown, 1) + wcov(ij, 1)) &
                                  *ne_ldown*(eta_dot(ij + t_ldown, 1) - eta_dot(ij, 1))
         END DO
      END DO
      ! add NH terms to dW, dPhi
      ! FIXME : TODO top and bottom
      DO l = ll_beginp1, ll_end ! inner interfaces only
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dPhi(ij, l) = dPhi(ij, l) - wflux(ij, l) &
                          *(geopot(ij, l + 1) - geopot(ij, l - 1))/(mass(ij, l - 1) + mass(ij, l))
         END DO
      END DO
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dW(ij, l + 1) = dW(ij, l + 1) + W_etadot(ij, l) ! update inner+top interfaces
            dW(ij, l) = dW(ij, l) - W_etadot(ij, l) ! update bottom+inner interfaces
         END DO
      END DO

      CALL trace_end("compute_caldyn_vert_nh")

   END SUBROUTINE compute_caldyn_vert_NH_manual

   SUBROUTINE compute_caldyn_vert_NH_hex(mass, geopot, W, wflux, W_etadot, du, dPhi, dW)
      USE dimensions ! iim, jjm, offsets ...
      USE geometry ! signs
      USE grid_param, ONLY: llm, nqdyn
      USE trace, ONLY: trace_start, trace_end
      USE omp_para, ONLY: ll_begin, ll_beginp1, ll_end, ll_endp1, ll_endm1
      REAL(rstd), INTENT(IN) :: mass(iim*jjm, llm)
      REAL(rstd), INTENT(IN) :: geopot(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN) :: W(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN) :: wflux(iim*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: du(iim*3*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: dPhi(iim*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: dW(iim*jjm, llm + 1)
      REAL(rstd) :: W_etadot(iim*jjm, llm) ! vertical flux of vertical momentum
      ! local arrays
      REAL(rstd) :: eta_dot(iim*jjm, llm) ! eta_dot in full layers
      REAL(rstd) :: wcov(iim*jjm, llm) ! covariant vertical momentum in full layers
      ! indices and temporary values
      INTEGER    :: ij, l
      REAL(rstd) :: wflux_ij, w_ij

      CALL trace_start("compute_caldyn_vert_nh")

      !$OMP BARRIER
      IF (ll_begin == 1) THEN
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            w_ij = (W(ij, 1) + .5*W(ij, 1 + 1))/mass(ij, 1)
            wflux_ij = .5*(wflux(ij, 1) + wflux(ij, 1 + 1))
            W_etadot(ij, 1) = wflux_ij*w_ij
            eta_dot(ij, 1) = wflux_ij/mass(ij, 1)
            wcov(ij, 1) = w_ij*(geopot(ij, 1 + 1) - geopot(ij, 1))
         END DO
      END IF
      DO l = ll_beginp1, ll_endm1
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            w_ij = .5*(W(ij, l) + W(ij, l + 1))/mass(ij, l)
            wflux_ij = .5*(wflux(ij, l) + wflux(ij, l + 1))
            W_etadot(ij, l) = wflux_ij*w_ij
            eta_dot(ij, l) = wflux_ij/mass(ij, l)
            wcov(ij, l) = w_ij*(geopot(ij, l + 1) - geopot(ij, l))
         END DO
      END DO
      IF (ll_end == llm) THEN
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            w_ij = (.5*W(ij, llm) + W(ij, llm + 1))/mass(ij, llm)
            wflux_ij = .5*(wflux(ij, llm) + wflux(ij, llm + 1))
            W_etadot(ij, llm) = wflux_ij*w_ij
            eta_dot(ij, llm) = wflux_ij/mass(ij, llm)
            wcov(ij, llm) = w_ij*(geopot(ij, llm + 1) - geopot(ij, llm))
         END DO
      END IF
      ! add NH term 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) - &
             &                   .5*(wcov(ij, l) + wcov(ij + t_right, l))*ne_right*(eta_dot(ij + t_right, l) - eta_dot(ij, l))
            du(ij + u_lup, l) = du(ij + u_lup, l) - &
            &                   .5*(wcov(ij, l) + wcov(ij + t_lup, l))*ne_lup*(eta_dot(ij + t_lup, l) - eta_dot(ij, l))
            du(ij + u_ldown, l) = du(ij + u_ldown, l) - &
           &                     .5*(wcov(ij, l) + wcov(ij + t_ldown, l))*ne_ldown*(eta_dot(ij + t_ldown, l) - eta_dot(ij, l))
         END DO
      END DO
      ! add NH terms to dW, dPhi
      ! FIXME : TODO top and bottom
      DO l = ll_beginp1, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dPhi(ij, l) = dPhi(ij, l) - wflux(ij, l)*(geopot(ij, l + 1) - geopot(ij, l - 1))/(mass(ij, l - 1) + mass(ij, l))
         END DO
      END DO
      ! We need a barrier here because we compute W_etadot above and do a vertical difference below
      !$OMP BARRIER
      IF (ll_begin == 1) THEN
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dW(ij, 1) = dW(ij, 1) - W_etadot(ij, 1)
         END DO
      END IF
      DO l = ll_beginp1, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dW(ij, l) = dW(ij, l) + W_etadot(ij, l - 1) - W_etadot(ij, l)
         END DO
      END DO
      IF (ll_endp1 == llm + 1) THEN
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dW(ij, llm + 1) = dW(ij, llm + 1) + W_etadot(ij, llm + 1 - 1)
         END DO
      END IF
      !$OMP BARRIER

      CALL trace_end("compute_caldyn_vert_nh")
   END SUBROUTINE compute_caldyn_vert_NH_hex

   SUBROUTINE compute_caldyn_vert_NH_unst(mass, geopot, W, wflux, eta_dot, wcov, W_etadot, du, dPhi, dW)
      USE grid_param
      USE ISO_C_BINDING, only: C_DOUBLE, C_FLOAT
      USE data_unstructured_mod, ONLY: enter_trace, exit_trace, &
                                       id_vert_NH, left, right
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)   :: mass, eta_dot, wcov, W_etadot ! IN, BUF*3
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: geopot, W, wflux, dPhi, dW ! IN*3, INOUT*2
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)      :: du ! INOUT
      INTEGER ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown
      REAL(C_DOUBLE) :: w_ij, wflux_ij
      CALL enter_trace(id_vert_NH, 8*llm*((6)*primal_num + (0)*dual_num + (1)*edge_num))

      !--------------------------------------------------------------------------
      !---------------------------- caldyn_vert_NH ----------------------------------
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         l = 1
         w_ij = (W(l, ij) + .5*W(l + 1, ij))/mass(l, ij)
         wflux_ij = .5*(wflux(l, ij) + wflux(l + 1, ij))
         W_etadot(l, ij) = wflux_ij*w_ij
         eta_dot(l, ij) = wflux_ij/mass(l, ij)
         wcov(l, ij) = w_ij*(geopot(l + 1, ij) - geopot(l, ij))
         !DIR$ SIMD
         DO l = 2, llm - 1
            w_ij = .5*(W(l, ij) + W(l + 1, ij))/mass(l, ij)
            wflux_ij = .5*(wflux(l, ij) + wflux(l + 1, ij))
            W_etadot(l, ij) = wflux_ij*w_ij
            eta_dot(l, ij) = wflux_ij/mass(l, ij)
            wcov(l, ij) = w_ij*(geopot(l + 1, ij) - geopot(l, ij))
         END DO
         l = llm
         w_ij = (.5*W(l, ij) + W(l + 1, ij))/mass(l, ij)
         wflux_ij = .5*(wflux(l, ij) + wflux(l + 1, ij))
         W_etadot(l, ij) = wflux_ij*w_ij
         eta_dot(l, ij) = wflux_ij/mass(l, ij)
         wcov(l, ij) = w_ij*(geopot(l + 1, ij) - geopot(l, ij))
      END DO
      !$OMP END DO
      ! add NH term to du
      !$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) - .5*(wcov(l, ij_left) + wcov(l, ij_right))*1.*(eta_dot(l, ij_right) - eta_dot(l, ij_left))
         END DO
      END DO
      !$OMP END DO
      ! add NH terms to dW, dPhi
      ! FIXME : TODO top and bottom
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         !DIR$ SIMD
         DO l = 2, llm
            dPhi(l, ij) = dPhi(l, ij) - wflux(l, ij)*(geopot(l + 1, ij) - geopot(l - 1, ij))/(mass(l - 1, ij) + mass(l, ij))
         END DO
      END DO
      !$OMP END DO
      ! We need a barrier here because we compute W_etadot above and do a vertical difference below
      !$OMP BARRIER
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         l = 1
         dW(l, ij) = dW(l, ij) - W_etadot(l, ij)
         !DIR$ SIMD
         DO l = 2, llm
            dW(l, ij) = dW(l, ij) + W_etadot(l - 1, ij) - W_etadot(l, ij)
         END DO
         l = llm + 1
         dW(l, ij) = dW(l, ij) + W_etadot(l - 1, ij)
      END DO
      !$OMP END DO
      !---------------------------- caldyn_vert_NH ----------------------------------
      !--------------------------------------------------------------------------
      CALL exit_trace()
   END SUBROUTINE compute_caldyn_vert_NH_unst

END SUBMODULE compute_caldyn_vert_NH_mod
