SUBMODULE(compute_caldyn_mod) compute_caldyn_horiz_mod

   IMPLICIT NONE

CONTAINS

   MODULE SUBROUTINE compute_caldyn_horiz(u, rhodz, qu, theta, pk, geopot, hflux, convm, dtheta_rhodz, du)
      USE prec, ONLY: rstd
      REAL(rstd), INTENT(IN)  :: u(:, :)    ! prognostic "velocity"
      REAL(rstd), INTENT(IN)  :: rhodz(:, :)
      REAL(rstd), INTENT(IN)  :: qu(:, :)
      REAL(rstd), INTENT(IN)  :: theta(:, :)  ! potential temperature
      REAL(rstd), INTENT(INOUT) :: pk(:, :) ! Exner function
      REAL(rstd), INTENT(IN)  :: geopot(:, :)    ! geopotential

      REAL(rstd), INTENT(OUT) :: hflux(:, :) ! hflux in kg/s
      REAL(rstd), INTENT(OUT) :: convm(:, :)  ! mass flux convergence
      REAL(rstd), INTENT(OUT) :: dtheta_rhodz(:, :)
      REAL(rstd), INTENT(OUT) :: du(:, :)

      CALL compute_caldyn_horiz_hex(u, rhodz, qu, theta, pk, geopot, hflux, convm, dtheta_rhodz, du)
   END SUBROUTINE

   SUBROUTINE compute_caldyn_horiz_hex(u, rhodz, qu, theta, pk, geopot, hflux, convm, dtheta_rhodz, du)
      USE icosa
      USE disvert_mod
      USE trace
      USE omp_para
      USE caldyn_vars_mod
      REAL(rstd), INTENT(IN)  :: u(iim*3*jjm, llm)    ! prognostic "velocity"
      REAL(rstd), INTENT(IN)  :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(IN)  :: qu(iim*3*jjm, llm)
      REAL(rstd), INTENT(IN)  :: theta(iim*jjm, llm)  ! potential temperature
      REAL(rstd), INTENT(INOUT) :: pk(iim*jjm, llm) ! Exner function
      REAL(rstd), INTENT(IN)  :: geopot(iim*jjm, llm + 1)    ! geopotential

      REAL(rstd), INTENT(OUT) :: hflux(iim*3*jjm, llm) ! hflux in kg/s
      REAL(rstd), INTENT(OUT) :: convm(iim*jjm, llm)  ! mass flux convergence
      REAL(rstd), INTENT(OUT) :: dtheta_rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: du(iim*3*jjm, llm)

      REAL(rstd) :: Ftheta(3*iim*jjm, llm) ! theta flux
      REAL(rstd) :: berni(iim*jjm, llm)  ! Bernoulli function
      REAL(rstd) :: uu_right, uu_lup, uu_ldown

      INTEGER :: ij, l, i, j
      REAL(rstd) :: uu, ww

      CALL trace_start("compute_caldyn_horiz")

      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            uu_right = 0.5*(rhodz(ij, l) + rhodz(ij + t_right, l))*u(ij + u_right, l)
            uu_lup = 0.5*(rhodz(ij, l) + rhodz(ij + t_lup, l))*u(ij + u_lup, l)
            uu_ldown = 0.5*(rhodz(ij, l) + rhodz(ij + t_ldown, l))*u(ij + u_ldown, l)
            uu_right = uu_right*le_de(ij + u_right)
            uu_lup = uu_lup*le_de(ij + u_lup)
            uu_ldown = uu_ldown*le_de(ij + u_ldown)
            hflux(ij + u_right, l) = uu_right
            hflux(ij + u_lup, l) = uu_lup
            hflux(ij + u_ldown, l) = uu_ldown
            Ftheta(ij + u_right, l) = 0.5*(theta(ij, l) + theta(ij + t_right, l))*hflux(ij + u_right, l)
            Ftheta(ij + u_lup, l) = 0.5*(theta(ij, l) + theta(ij + t_lup, l))*hflux(ij + u_lup, l)
            Ftheta(ij + u_ldown, l) = 0.5*(theta(ij, l) + theta(ij + t_ldown, l))*hflux(ij + u_ldown, l)
         END DO

         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            ! convm = -div(mass flux), sign convention as in Ringler et al. 2012, eq. 21
            convm(ij, l) = -1./Ai(ij)*(ne_right*hflux(ij + u_right, l) + &
                                       ne_rup*hflux(ij + u_rup, l) + &
                                       ne_lup*hflux(ij + u_lup, l) + &
                                       ne_left*hflux(ij + u_left, l) + &
                                       ne_ldown*hflux(ij + u_ldown, l) + &
                                       ne_rdown*hflux(ij + u_rdown, l))

            ! dtheta_rhodz =  -div(flux.theta)
            dtheta_rhodz(ij, l) = -1./Ai(ij)*(ne_right*Ftheta(ij + u_right, l) + &
                                              ne_rup*Ftheta(ij + u_rup, l) + &
                                              ne_lup*Ftheta(ij + u_lup, l) + &
                                              ne_left*Ftheta(ij + u_left, l) + &
                                              ne_ldown*Ftheta(ij + u_ldown, l) + &
                                              ne_rdown*Ftheta(ij + u_rdown, l))
         END DO

      END DO

      SELECT CASE (caldyn_conserv)
      CASE (conserv_energy) ! energy-conserving TRiSK

         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end

               uu = wee(ij + u_right, 1)*hflux(ij + u_rup, l)*(qu(ij + u_right, l) + qu(ij + u_rup, l)) + &
                    wee(ij + u_right, 2)*hflux(ij + u_lup, l)*(qu(ij + u_right, l) + qu(ij + u_lup, l)) + &
                    wee(ij + u_right, 3)*hflux(ij + u_left, l)*(qu(ij + u_right, l) + qu(ij + u_left, l)) + &
                    wee(ij + u_right, 4)*hflux(ij + u_ldown, l)*(qu(ij + u_right, l) + qu(ij + u_ldown, l)) + &
                    wee(ij + u_right, 5)*hflux(ij + u_rdown, l)*(qu(ij + u_right, l) + qu(ij + u_rdown, l)) + &
                    wee(ij + u_right, 6)*hflux(ij + t_right + u_ldown, l)*(qu(ij + u_right, l) + qu(ij + t_right + u_ldown, l)) + &
                    wee(ij + u_right, 7)*hflux(ij + t_right + u_rdown, l)*(qu(ij + u_right, l) + qu(ij + t_right + u_rdown, l)) + &
                    wee(ij + u_right, 8)*hflux(ij + t_right + u_right, l)*(qu(ij + u_right, l) + qu(ij + t_right + u_right, l)) + &
                    wee(ij + u_right, 9)*hflux(ij + t_right + u_rup, l)*(qu(ij + u_right, l) + qu(ij + t_right + u_rup, l)) + &
                    wee(ij + u_right, 10)*hflux(ij + t_right + u_lup, l)*(qu(ij + u_right, l) + qu(ij + t_right + u_lup, l))
               du(ij + u_right, l) = .5*uu

               uu = wee(ij + u_lup, 1)*hflux(ij + u_left, l)*(qu(ij + u_lup, l) + qu(ij + u_left, l)) + &
                    wee(ij + u_lup, 2)*hflux(ij + u_ldown, l)*(qu(ij + u_lup, l) + qu(ij + u_ldown, l)) + &
                    wee(ij + u_lup, 3)*hflux(ij + u_rdown, l)*(qu(ij + u_lup, l) + qu(ij + u_rdown, l)) + &
                    wee(ij + u_lup, 4)*hflux(ij + u_right, l)*(qu(ij + u_lup, l) + qu(ij + u_right, l)) + &
                    wee(ij + u_lup, 5)*hflux(ij + u_rup, l)*(qu(ij + u_lup, l) + qu(ij + u_rup, l)) + &
                    wee(ij + u_lup, 6)*hflux(ij + t_lup + u_right, l)*(qu(ij + u_lup, l) + qu(ij + t_lup + u_right, l)) + &
                    wee(ij + u_lup, 7)*hflux(ij + t_lup + u_rup, l)*(qu(ij + u_lup, l) + qu(ij + t_lup + u_rup, l)) + &
                    wee(ij + u_lup, 8)*hflux(ij + t_lup + u_lup, l)*(qu(ij + u_lup, l) + qu(ij + t_lup + u_lup, l)) + &
                    wee(ij + u_lup, 9)*hflux(ij + t_lup + u_left, l)*(qu(ij + u_lup, l) + qu(ij + t_lup + u_left, l)) + &
                    wee(ij + u_lup, 10)*hflux(ij + t_lup + u_ldown, l)*(qu(ij + u_lup, l) + qu(ij + t_lup + u_ldown, l))
               du(ij + u_lup, l) = .5*uu

               uu = wee(ij + u_ldown, 1)*hflux(ij + u_rdown, l)*(qu(ij + u_ldown, l) + qu(ij + u_rdown, l)) + &
                    wee(ij + u_ldown, 2)*hflux(ij + u_right, l)*(qu(ij + u_ldown, l) + qu(ij + u_right, l)) + &
                    wee(ij + u_ldown, 3)*hflux(ij + u_rup, l)*(qu(ij + u_ldown, l) + qu(ij + u_rup, l)) + &
                    wee(ij + u_ldown, 4)*hflux(ij + u_lup, l)*(qu(ij + u_ldown, l) + qu(ij + u_lup, l)) + &
                    wee(ij + u_ldown, 5)*hflux(ij + u_left, l)*(qu(ij + u_ldown, l) + qu(ij + u_left, l)) + &
                    wee(ij + u_ldown, 6)*hflux(ij + t_ldown + u_lup, l)*(qu(ij + u_ldown, l) + qu(ij + t_ldown + u_lup, l)) + &
                    wee(ij + u_ldown, 7)*hflux(ij + t_ldown + u_left, l)*(qu(ij + u_ldown, l) + qu(ij + t_ldown + u_left, l)) + &
                    wee(ij + u_ldown, 8)*hflux(ij + t_ldown + u_ldown, l)*(qu(ij + u_ldown, l) + qu(ij + t_ldown + u_ldown, l)) + &
                    wee(ij + u_ldown, 9)*hflux(ij + t_ldown + u_rdown, l)*(qu(ij + u_ldown, l) + qu(ij + t_ldown + u_rdown, l)) + &
                    wee(ij + u_ldown, 10)*hflux(ij + t_ldown + u_right, l)*(qu(ij + u_ldown, l) + qu(ij + t_ldown + u_right, l))
               du(ij + u_ldown, l) = .5*uu

            END DO
         END DO

      CASE (conserv_enstrophy) ! enstrophy-conserving TRiSK

         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end

               uu = wee(ij + u_right, 1)*hflux(ij + u_rup, l) + &
                    wee(ij + u_right, 2)*hflux(ij + u_lup, l) + &
                    wee(ij + u_right, 3)*hflux(ij + u_left, l) + &
                    wee(ij + u_right, 4)*hflux(ij + u_ldown, l) + &
                    wee(ij + u_right, 5)*hflux(ij + u_rdown, l) + &
                    wee(ij + u_right, 6)*hflux(ij + t_right + u_ldown, l) + &
                    wee(ij + u_right, 7)*hflux(ij + t_right + u_rdown, l) + &
                    wee(ij + u_right, 8)*hflux(ij + t_right + u_right, l) + &
                    wee(ij + u_right, 9)*hflux(ij + t_right + u_rup, l) + &
                    wee(ij + u_right, 10)*hflux(ij + t_right + u_lup, l)
               du(ij + u_right, l) = qu(ij + u_right, l)*uu

               uu = wee(ij + u_lup, 1)*hflux(ij + u_left, l) + &
                    wee(ij + u_lup, 2)*hflux(ij + u_ldown, l) + &
                    wee(ij + u_lup, 3)*hflux(ij + u_rdown, l) + &
                    wee(ij + u_lup, 4)*hflux(ij + u_right, l) + &
                    wee(ij + u_lup, 5)*hflux(ij + u_rup, l) + &
                    wee(ij + u_lup, 6)*hflux(ij + t_lup + u_right, l) + &
                    wee(ij + u_lup, 7)*hflux(ij + t_lup + u_rup, l) + &
                    wee(ij + u_lup, 8)*hflux(ij + t_lup + u_lup, l) + &
                    wee(ij + u_lup, 9)*hflux(ij + t_lup + u_left, l) + &
                    wee(ij + u_lup, 10)*hflux(ij + t_lup + u_ldown, l)
               du(ij + u_lup, l) = qu(ij + u_lup, l)*uu

               uu = wee(ij + u_ldown, 1)*hflux(ij + u_rdown, l) + &
                    wee(ij + u_ldown, 2)*hflux(ij + u_right, l) + &
                    wee(ij + u_ldown, 3)*hflux(ij + u_rup, l) + &
                    wee(ij + u_ldown, 4)*hflux(ij + u_lup, l) + &
                    wee(ij + u_ldown, 5)*hflux(ij + u_left, l) + &
                    wee(ij + u_ldown, 6)*hflux(ij + t_ldown + u_lup, l) + &
                    wee(ij + u_ldown, 7)*hflux(ij + t_ldown + u_left, l) + &
                    wee(ij + u_ldown, 8)*hflux(ij + t_ldown + u_ldown, l) + &
                    wee(ij + u_ldown, 9)*hflux(ij + t_ldown + u_rdown, l) + &
                    wee(ij + u_ldown, 10)*hflux(ij + t_ldown + u_right, l)
               du(ij + u_ldown, l) = qu(ij + u_ldown, l)*uu

            END DO
         END DO

      CASE DEFAULT
         PRINT *, 'Invalid value of caldyn_conserv : ', caldyn_conserv
         STOP
      END SELECT

      IF (boussinesq) THEN
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               berni(ij, l) = pk(ij, l) + 1/(4*Ai(ij))*( &
                              le_de(ij + u_right)*u(ij + u_right, l)**2 + &
                              le_de(ij + u_rup)*u(ij + u_rup, l)**2 + &
                              le_de(ij + u_lup)*u(ij + u_lup, l)**2 + &
                              le_de(ij + u_left)*u(ij + u_left, l)**2 + &
                              le_de(ij + u_ldown)*u(ij + u_ldown, l)**2 + &
                              le_de(ij + u_rdown)*u(ij + u_rdown, l)**2)
               ! from now on pk contains the vertically-averaged geopotential
               pk(ij, l) = .5*(geopot(ij, l) + geopot(ij, l + 1))
            END DO
         END DO

      ELSE ! compressible

         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               berni(ij, l) = .5*(geopot(ij, l) + geopot(ij, l + 1)) &
                              + 1/(4*Ai(ij))*( &
                              le_de(ij + u_right)*u(ij + u_right, l)**2 + &
                              le_de(ij + u_rup)*u(ij + u_rup, l)**2 + &
                              le_de(ij + u_lup)*u(ij + u_lup, l)**2 + &
                              le_de(ij + u_left)*u(ij + u_left, l)**2 + &
                              le_de(ij + u_ldown)*u(ij + u_ldown, l)**2 + &
                              le_de(ij + u_rdown)*u(ij + u_rdown, l)**2)
            END DO
         END DO

      END IF ! Boussinesq/compressible

      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end

            du(ij + u_right, l) = du(ij + u_right, l) + ( &
                                  0.5*(theta(ij, l) + theta(ij + t_right, l)) &
                                  *(ne_right*pk(ij, l) + ne_left*pk(ij + t_right, l)) &
                                  + ne_right*berni(ij, l) + ne_left*berni(ij + t_right, l))

            du(ij + u_lup, l) = du(ij + u_lup, l) + ( &
                                0.5*(theta(ij, l) + theta(ij + t_lup, l)) &
                                *(ne_lup*pk(ij, l) + ne_rdown*pk(ij + t_lup, l)) &
                                + ne_lup*berni(ij, l) + ne_rdown*berni(ij + t_lup, l))

            du(ij + u_ldown, l) = du(ij + u_ldown, l) + ( &
                                  0.5*(theta(ij, l) + theta(ij + t_ldown, l)) &
                                  *(ne_ldown*pk(ij, l) + ne_rup*pk(ij + t_ldown, l)) &
                                  + ne_ldown*berni(ij, l) + ne_rup*berni(ij + t_ldown, l))

         END DO
      END DO

      CALL trace_end("compute_caldyn_horiz")

   END SUBROUTINE compute_caldyn_horiz_hex

END SUBMODULE
