MODULE accumulate_fluxes_mod
   USE prec, ONLY: rstd
   IMPLICIT NONE
   PRIVATE

   PUBLIC :: accumulate_fluxes_hex, accumulate_fluxes_unst

CONTAINS

   SUBROUTINE accumulate_fluxes_hex(hflux, wflux, hfluxt, wfluxt, tau, fluxt_zero)
      USE dimensions, ONLY: iim, jjm, ij_begin, ij_end, ij_begin_ext, ij_end_ext, u_right, u_lup, u_ldown
      USE omp_para, ONLY: ll_begin, ll_end, ll_endp1
      USE disvert_mod, ONLY: llm, caldyn_eta, eta_mass
      REAL(rstd), INTENT(IN)    :: hflux(3*iim*jjm, llm), wflux(iim*jjm, llm + 1)
      REAL(rstd), INTENT(INOUT) :: hfluxt(3*iim*jjm, llm), wfluxt(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN) :: tau
      LOGICAL, INTENT(INOUT) :: fluxt_zero
      INTEGER :: l, ij

      !$acc data present(hflux(:,:), wflux(:,:), hfluxt(:,:), wfluxt(:,:)) async

      IF (fluxt_zero) THEN

         fluxt_zero = .FALSE.
         !$acc parallel loop async
         DO l = ll_begin, ll_end
            !$acc loop
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               hfluxt(ij + u_right, l) = tau*hflux(ij + u_right, l)
               hfluxt(ij + u_lup, l) = tau*hflux(ij + u_lup, l)
               hfluxt(ij + u_ldown, l) = tau*hflux(ij + u_ldown, l)
            END DO
         END DO

         IF (caldyn_eta == eta_mass) THEN ! no need for vertical fluxes if eta_lag
            !$acc parallel loop async
            DO l = ll_begin, ll_endp1
               !$acc loop
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  wfluxt(ij, l) = tau*wflux(ij, l)
               END DO
            END DO
         END IF

      ELSE
         !$acc parallel loop async
         DO l = ll_begin, ll_end
            !$acc loop
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               hfluxt(ij + u_right, l) = hfluxt(ij + u_right, l) + tau*hflux(ij + u_right, l)
               hfluxt(ij + u_lup, l) = hfluxt(ij + u_lup, l) + tau*hflux(ij + u_lup, l)
               hfluxt(ij + u_ldown, l) = hfluxt(ij + u_ldown, l) + tau*hflux(ij + u_ldown, l)
            END DO
         END DO

         IF (caldyn_eta == eta_mass) THEN ! no need for vertical fluxes if eta_lag
            !$acc parallel loop async
            DO l = ll_begin, ll_endp1
               !$acc loop
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  wfluxt(ij, l) = wfluxt(ij, l) + tau*wflux(ij, l)
               END DO
            END DO
         END IF
      END IF
      !$acc end data

   END SUBROUTINE accumulate_fluxes_hex

   SUBROUTINE accumulate_fluxes_unst(hflux, wflux, hfluxt, wfluxt, tau, fluxt_zero)
      USE grid_param, ONLY: primal_num, edge_num
      USE disvert_mod, ONLY: llm, caldyn_eta, eta_mass
      INTEGER :: l, ij, edge
      REAL(rstd), INTENT(IN)    :: hflux(llm, edge_num), wflux(llm + 1, primal_num)
      REAL(rstd), INTENT(INOUT) :: hfluxt(llm, edge_num), wfluxt(llm + 1, primal_num)
      REAL(rstd), INTENT(IN) :: tau
      LOGICAL, INTENT(INOUT) :: fluxt_zero

      IF (fluxt_zero) THEN
         fluxt_zero = .FALSE.
         !$OMP DO SCHEDULE(STATIC)
         DO edge = 1, edge_num
            !DIR$ SIMD
            DO l = 1, llm
               hfluxt(l, edge) = tau*hflux(l, edge)
            END DO
         END DO
         !$OMP END DO
         IF (caldyn_eta == eta_mass) THEN ! no need for vertical fluxes if eta_lag
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm + 1
                  wfluxt(l, ij) = tau*wflux(l, ij)
               END DO
            END DO
            !$OMP END DO
         END IF
      ELSE
         !$OMP DO SCHEDULE(STATIC)
         DO edge = 1, edge_num
            !DIR$ SIMD
            DO l = 1, llm
               hfluxt(l, edge) = hfluxt(l, edge) + tau*hflux(l, edge)
            END DO
         END DO
         !$OMP END DO
         IF (caldyn_eta == eta_mass) THEN ! no need for vertical fluxes if eta_lag
            !$OMP DO SCHEDULE(STATIC)
            DO ij = 1, primal_num
               !DIR$ SIMD
               DO l = 1, llm + 1
                  wfluxt(l, ij) = wfluxt(l, ij) + tau*wflux(l, ij)
               END DO
            END DO
            !$OMP END DO
         END IF
      END IF
   END SUBROUTINE accumulate_fluxes_unst

END MODULE accumulate_fluxes_mod
