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

CONTAINS

   MODULE SUBROUTINE compute_caldyn_fast(tau, theta, geopot, pk, berni, du, u)
      USE trace
      USE grid_param, ONLY: grid_type, grid_ico
      REAL(rstd), INTENT(IN)    :: tau, theta(:, :, :), geopot(:, :)
      REAL(rstd), INTENT(INOUT) :: pk(:, :, :), berni(:, :), du(:, :), u(:, :)

      CALL trace_start("compute_caldyn_fast")
      IF (grid_type == grid_ico) THEN
         CALL compute_caldyn_fast_manual(tau, theta, geopot, pk, berni, du, u)
      ELSE
         CALL compute_caldyn_fast_unst_manual(tau, theta, geopot, pk(:, :, 1), berni, du, u)
      END IF
      CALL trace_end("compute_caldyn_fast")
   END SUBROUTINE compute_caldyn_fast

   SUBROUTINE compute_caldyn_fast_manual(tau, theta, geopot, pk, berni, du, u)
      USE icosa
      USE omp_para, ONLY: ll_begin, ll_end
      REAL(rstd), INTENT(IN)    :: tau                ! "solve" u-tau*du/dt = rhs
      REAL(rstd), INTENT(IN)    :: theta(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(IN)    :: geopot(iim*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: pk(iim*jjm, llm, nqdyn) ! Exner functions
      REAL(rstd), INTENT(INOUT) :: berni(iim*jjm, llm)  ! partial Bernoulli function
      REAL(rstd), INTENT(INOUT) :: du(iim*3*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: u(iim*3*jjm, llm)   ! INOUT if tau>0

      INTEGER :: ij, l
      REAL(rstd) :: phi_ik, due, due_right, due_lup, due_ldown

      !$acc data async &
      !$acc present(theta(:,:,:), geopot(:,:), pk(:,:), berni(:,:), du(:,:), u(:,:))

      ! Compute Bernoulli term
      IF (boussinesq) THEN
         !$acc parallel loop collapse(2) async
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               berni(ij, l) = pk(ij, l, 1)
               ! from now on pk contains the vertically-averaged geopotential
               pk(ij, l, 1) = .5*(geopot(ij, l) + geopot(ij, l + 1))
            END DO
         END DO
      ELSE ! compressible

         SELECT CASE (nqdyn)
         CASE (2)
            !$acc parallel loop collapse(2) async
            DO l = ll_begin, ll_end
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  phi_ik = .5*(geopot(ij, l) + geopot(ij, l + 1))
                  berni(ij, l) = phi_ik + (berni(ij, l) &
                  &                        - pk(ij, l, 1)*theta(ij, l, 1) &
                  &                        - pk(ij, l, 2)*theta(ij, l, 2))
               END DO
            END DO

         CASE (1)
            !$acc parallel loop collapse(2) async
            DO l = ll_begin, ll_end
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  phi_ik = .5*(geopot(ij, l) + geopot(ij, l + 1))
                  berni(ij, l) = phi_ik + (berni(ij, l) &
                  &                        - pk(ij, l, 1)*theta(ij, l, 1))
               END DO
            END DO
         END SELECT

      END IF ! Boussinesq/compressible

! u:=u+tau*du, du = -grad(B)-theta.grad(pi)
      SELECT CASE (nqdyn)
      CASE (2) ! binary fluid
         !$acc parallel loop collapse(2) async
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               due_right = berni(ij + t_right, l) - berni(ij, l) &
                           + 0.5*(theta(ij, l, 1) + theta(ij + t_right, l, 1)) &
                           *(pk(ij + t_right, l, 1) - pk(ij, l, 1)) &
                           + 0.5*(theta(ij, l, 2) + theta(ij + t_right, l, 2)) &
                           *(pk(ij + t_right, l, 2) - pk(ij, l, 2))

               due_lup = berni(ij + t_lup, l) - berni(ij, l) &
                         + 0.5*(theta(ij, l, 1) + theta(ij + t_lup, l, 1)) &
                         *(pk(ij + t_lup, l, 1) - pk(ij, l, 1)) &
                         + 0.5*(theta(ij, l, 2) + theta(ij + t_lup, l, 2)) &
                         *(pk(ij + t_lup, l, 2) - pk(ij, l, 2))

               due_ldown = berni(ij + t_ldown, l) - berni(ij, l) &
                           + 0.5*(theta(ij, l, 1) + theta(ij + t_ldown, l, 1)) &
                           *(pk(ij + t_ldown, l, 1) - pk(ij, l, 1)) &
                           + 0.5*(theta(ij, l, 2) + theta(ij + t_ldown, l, 2)) &
                           *(pk(ij + t_ldown, l, 2) - pk(ij, l, 2))

               du(ij + u_right, l) = du(ij + u_right, l) - ne_right*due_right
               du(ij + u_lup, l) = du(ij + u_lup, l) - ne_lup*due_lup
               du(ij + u_ldown, l) = du(ij + u_ldown, l) - ne_ldown*due_ldown
               u(ij + u_right, l) = u(ij + u_right, l) + tau*du(ij + u_right, l)
               u(ij + u_lup, l) = u(ij + u_lup, l) + tau*du(ij + u_lup, l)
               u(ij + u_ldown, l) = u(ij + u_ldown, l) + tau*du(ij + u_ldown, l)
            END DO
         END DO
      CASE (1)
         !$acc parallel loop collapse(2) async
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               due_right = 0.5*(theta(ij, l, 1) + theta(ij + t_right, l, 1)) &
                           *(pk(ij + t_right, l, 1) - pk(ij, l, 1)) &
                           + berni(ij + t_right, l) - berni(ij, l)
               due_lup = 0.5*(theta(ij, l, 1) + theta(ij + t_lup, l, 1)) &
                         *(pk(ij + t_lup, l, 1) - pk(ij, l, 1)) &
                         + berni(ij + t_lup, l) - berni(ij, l)
               due_ldown = 0.5*(theta(ij, l, 1) + theta(ij + t_ldown, l, 1)) &
                           *(pk(ij + t_ldown, l, 1) - pk(ij, l, 1)) &
                           + berni(ij + t_ldown, l) - berni(ij, l)
               du(ij + u_right, l) = du(ij + u_right, l) - ne_right*due_right
               du(ij + u_lup, l) = du(ij + u_lup, l) - ne_lup*due_lup
               du(ij + u_ldown, l) = du(ij + u_ldown, l) - ne_ldown*due_ldown
               u(ij + u_right, l) = u(ij + u_right, l) + tau*du(ij + u_right, l)
               u(ij + u_lup, l) = u(ij + u_lup, l) + tau*du(ij + u_lup, l)
               u(ij + u_ldown, l) = u(ij + u_ldown, l) + tau*du(ij + u_ldown, l)
            END DO
         END DO
      END SELECT

      !$acc end data
   END SUBROUTINE compute_caldyn_fast_manual

   SUBROUTINE compute_caldyn_fast_unst_manual(tau, theta, geopot, pk, berni, du, u)
      USE grid_param, ONLY: llm, primal_num, edge_num, nqdyn
      USE data_unstructured_mod, ONLY: left, right
      USE earth_const, ONLY: boussinesq
      REAL(rstd), INTENT(IN)    :: tau                ! "solve" u-tau*du/dt = rhs
      REAL(rstd), INTENT(IN)    :: theta(llm, primal_num, nqdyn)
      REAL(rstd), INTENT(IN)    :: geopot(llm + 1, primal_num)
      REAL(rstd), INTENT(INOUT) :: pk(llm, primal_num, nqdyn) ! Exner functions
      REAL(rstd), INTENT(INOUT) :: berni(llm, primal_num)  ! partial Bernoulli function
      REAL(rstd), INTENT(INOUT) :: du(llm, edge_num)
      REAL(rstd), INTENT(INOUT) :: u(llm, edge_num)        ! INOUT if tau>0

      INTEGER :: ij, l, edge, ij_left, ij_right
      REAL(rstd) :: phi_ik, due

      ! Compute Bernoulli term
      IF (boussinesq) THEN
         DO ij = 1, primal_num
            !DIR$ SIMD
            DO l = 1, llm
               berni(l, ij) = pk(l, ij, 1)
               ! from now on pk contains the vertically-averaged geopotential
               pk(l, ij, 1) = .5*(geopot(l, ij) + geopot(l + 1, ij))
            END DO
         END DO
      ELSE ! compressible

         SELECT CASE (nqdyn)
         CASE (2) ! binary fluid
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm
                  phi_ik = .5*(geopot(l, ij) + geopot(l + 1, ij))
                  berni(l, ij) = phi_ik + (berni(l, ij) &
                  &                        - pk(l, ij, 1)*theta(l, ij, 1) &
                  &                        - pk(l, ij, 2)*theta(l, ij, 2))
               END DO
            END DO

         CASE (1) ! single-component fluid
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm
                  phi_ik = .5*(geopot(l, ij) + geopot(l + 1, ij))
                  berni(l, ij) = phi_ik + (berni(l, ij) &
                  &                        - pk(l, ij, 1)*theta(l, ij, 1))
               END DO
            END DO
         END SELECT

      END IF ! Boussinesq/compressible

      ! u:=u+tau*du, du = -grad(B)-theta.grad(pi)
      SELECT CASE (nqdyn)
      CASE (2) ! binary fluid
         DO edge = 1, edge_num
            ij_left = left(edge)
            ij_right = right(edge)
            !DIR$ SIMD
            DO l = 1, llm
               due = (berni(l, ij_right) - berni(l, ij_left))            &
               &     + .5*(theta(l, ij_right, 1) + theta(l, ij_left, 1)) &
               &         *(pk(l, ij_right, 1) - pk(l, ij_left, 1))       &
               &     + .5*(theta(l, ij_right, 2) + theta(l, ij_left, 2)) &
               &         *(pk(l, ij_right, 2) - pk(l, ij_left, 2))
               du(l, edge) = du(l, edge) - due
               u(l, edge) = u(l, edge) + tau*du(l, edge)
            END DO
         END DO
      CASE (1) ! single-component fluid
         DO edge = 1, edge_num
            ij_left = left(edge)
            ij_right = right(edge)
            !DIR$ SIMD
            DO l = 1, llm
               due = (berni(l, ij_right) - berni(l, ij_left)) &
               &     + .5*(theta(l, ij_right, 1) + theta(l, ij_left, 1)) &
               &         *(pk(l, ij_right, 1) - pk(l, ij_left, 1))
               du(l, edge) = du(l, edge) - due
               u(l, edge) = u(l, edge) + tau*du(l, edge)
            END DO
         END DO
      END SELECT

   END SUBROUTINE compute_caldyn_fast_unst_manual

END SUBMODULE compute_caldyn_fast_mod
