! Original source is in jsrc

MODULE accumulate_fluxes_mod
  USE prec, ONLY : rstd
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: compute_accumulate_fluxes_hex, compute_accumulate_fluxes_unst
  
CONTAINS

#if BEGIN_DYSL

KERNEL(accumulate_fluxes)

IF(fluxt_zero) THEN
   fluxt_zero=.FALSE.
   FORALL_CELLS_EXT()
      ON_EDGES
        hfluxt(EDGE) = tau*hflux(EDGE)
      END_BLOCK
   END_BLOCK
      
   IF(caldyn_eta==eta_mass) THEN ! no need for vertical fluxes if eta_lag
     FORALL_CELLS('1','llm+1')
        ON_PRIMAL
          wfluxt(CELL) = tau*wflux(CELL)
        END_BLOCK
     END_BLOCK
  END IF

ELSE

   FORALL_CELLS_EXT()
      ON_EDGES
        hfluxt(EDGE) = hfluxt(EDGE) + tau*hflux(EDGE)
      END_BLOCK
   END_BLOCK
      
   IF(caldyn_eta==eta_mass) THEN ! no need for vertical fluxes if eta_lag
     FORALL_CELLS('1','llm+1')
        ON_PRIMAL
          wfluxt(CELL) = wfluxt(CELL) + tau*wflux(CELL)
        END_BLOCK
     END_BLOCK
  END IF

END IF

END_BLOCK
#endif END_DYSL

{% for grid in 'hex', 'unst' %}

  SUBROUTINE compute_accumulate_fluxes_{{grid}}(hflux,wflux, hfluxt,wfluxt, tau,fluxt_zero)
    REAL(rstd), INTENT(IN)    :: hflux(:,:), wflux(:,:)    ! horizontal and vertical mass fluxes
    REAL(rstd), INTENT(INOUT) :: hfluxt(:,:), wfluxt(:,:)  ! time-integrated fluxes
    REAL(rstd), INTENT(IN) :: tau                          ! time interval
    LOGICAL, INTENT(INOUT) :: fluxt_zero                   ! .TRUE. to reset fluxes
    CALL accumulate_fluxes_{{grid}}_(hflux,wflux, hfluxt,wfluxt, tau,fluxt_zero)
  END SUBROUTINE compute_accumulate_fluxes_{{grid}}

{% endfor %}
  
  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
#define SHAPE_EK  3*iim*jjm,llm
#define SHAPE_IL  iim*jjm,llm+1
    INTEGER :: l,ij
    REAL(rstd), INTENT(IN)    :: hflux(SHAPE_EK), wflux(SHAPE_IL)
    REAL(rstd), INTENT(INOUT) :: hfluxt(SHAPE_EK), wfluxt(SHAPE_IL)
    REAL(rstd), INTENT(IN) :: tau
    LOGICAL, INTENT(INOUT) :: fluxt_zero
#include "../kernels_hex/accumulate_fluxes.k90"
  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
#define SHAPE_EK llm,edge_num
#define SHAPE_IL llm+1,primal_num
    INTEGER :: l,ij,edge
    REAL(rstd), INTENT(IN)    :: hflux(SHAPE_EK), wflux(SHAPE_IL)
    REAL(rstd), INTENT(INOUT) :: hfluxt(SHAPE_EK), wfluxt(SHAPE_IL)
    REAL(rstd), INTENT(IN) :: tau
    LOGICAL, INTENT(INOUT) :: fluxt_zero
#include "../kernels_unst/accumulate_fluxes.k90"
  END SUBROUTINE accumulate_fluxes_unst_

END MODULE accumulate_fluxes_mod
