SUBMODULE(compute_transport_mod) compute_advect_horiz_mod

   IMPLICIT NONE

CONTAINS

   MODULE SUBROUTINE compute_advect_horiz(update_mass, diagflux_on, massp, fluxp, massq, fluxq, &
                                              qi, gradq3d, qfluxt, cc, qsat)
      USE grid_param, ONLY: grid_type, grid_ico
      USE geometry, ONLY: Ai, xyz_i
      LOGICAL, INTENT(IN)       :: update_mass, diagflux_on
      REAL(rstd), INTENT(IN)    :: fluxp(:, :)      ! parent mass flux
      REAL(rstd), INTENT(OUT)   :: fluxq(:, :)      ! tracer mass flux
      REAL(rstd), INTENT(INOUT) :: massp(:, :), massq(:, :) ! parent and tracer mass
      REAL(rstd), INTENT(INOUT) :: qi(:, :)         ! tracer concentration
      REAL(rstd), INTENT(IN)    :: gradq3d(:, :, :) ! tracer gradient
      REAL(rstd), INTENT(INOUT) :: qfluxt(:, :)     ! time-integrated tracer flux
      REAL(rstd), INTENT(IN)    :: cc(:, :, :)      ! barycenter of quadrilateral, where q is evaluated (1-point quadrature)
      REAL(rstd), INTENT(IN), OPTIONAL  :: qsat(:, :)   

      IF (grid_type == grid_ico) THEN
         CALL compute_advect_horiz_hex(update_mass, diagflux_on, massp, fluxp, massq, fluxq, &
                                       qi, gradq3d, qfluxt, cc, Ai, xyz_i, qsat)
      ELSE
         CALL compute_advect_horiz_unst(update_mass, diagflux_on, massp, fluxp, massq, fluxq, qi, qfluxt, qsat)
      END IF

   END SUBROUTINE

   SUBROUTINE compute_advect_horiz_unst(update_mass, diagflux_on, massp, fluxp, massq, fluxq, qi, qfluxt, qsat)
      USE data_unstructured_mod, ONLY: edge_num, left, right, primal_num, primal_deg, primal_edge, primal_ne
      USE grid_param, ONLY: llm, nqtot
      USE geometry, ONLY: Ai
      USE trace
      LOGICAL, INTENT(IN) :: update_mass, diagflux_on
      REAL(rstd), DIMENSION(llm, edge_num), INTENT(IN)      :: fluxp ! mass flux
      REAL(rstd), DIMENSION(llm, edge_num), INTENT(INOUT)   :: fluxq ! mass flux
      REAL(rstd), DIMENSION(llm, primal_num), INTENT(INOUT) :: massq, massp
      REAL(rstd), DIMENSION(llm, primal_num), INTENT(INOUT) :: qi
      REAL(rstd), DIMENSION(llm, edge_num), INTENT(INOUT) :: qfluxt    ! time-integrated tracer flux
      REAL(rstd), DIMENSION(llm, primal_num), OPTIONAL, INTENT(IN) :: qsat ! saturating water vapor in kg/kg, used only for vanleer_qsat scheme
 
      INTEGER :: ij, l, iq, iedge, edge, ivertex, vertex, ij_left, ij_right, ij_up, ij_down, itrisk, edge_trisk, kup, kdown
      REAL(rstd) :: sign1, sign2, sign3, sign4, sign5, sign6, le_de1, le_de2, le_de3, le_de4, le_de5, le_de6, &
                    de1, de2, de3, de4, de5, de6, le1, le2, le3, le4, le5, le6
      INTEGER :: edge1, edge2, edge3, edge4, edge5, edge6, ij_up1, ij_up2, ij_up3, ij_up4, ij_up5, ij_up6, &
                 ij_down1, ij_down2, ij_down3, ij_down4, ij_down5, ij_down6
      REAL(rstd), DIMENSION(llm, edge_num)    :: qflux
      REAL(rstd) :: qe, dmass, dq, newmass
      LOGICAL :: ok_qsat

      CALL trace_start("compute_advect_horiz")
      ok_qsat=PRESENT(qsat)
     
      !--------------------------------------------------------------------------
      !---------------------------- advect_horiz ----------------------------------
      ! evaluate tracer field at point cc using piecewise linear reconstruction
      ! q(cc)= q0 + gradq.(cc-xyz_i) with xi centroid of hexagon
      ! sign*fluxp>0 iff outgoing
      !$OMP DO SCHEDULE(STATIC)
      DO edge = 1, edge_num
         ij_left = left(edge)
         ij_right = right(edge)
         !DIR$ SIMD
         DO l = 1, llm
            IF (fluxp(l, edge) > 0.) THEN
               qe = qi(l, ij_left)
               qe = qe + (0.-0.)*0.
               qe = qe + (0.-0.)*0.
               qe = qe + (0.-0.)*0.
               IF (ok_qsat) qe = MIN(qsat(l,ij_left), qe)
            ELSE
               qe = qi(l, ij_right)
               qe = qe + (0.-0.)*0.
               qe = qe + (0.-0.)*0.
               qe = qe + (0.-0.)*0.
               IF (ok_qsat) qe = MIN(qsat(l,ij_right), qe)
            END IF

            fluxq(l, edge) = fluxp(l, edge)*qe
            IF (diagflux_on) qfluxt(l, edge) = qfluxt(l, edge) + fluxq(l, edge)
         END DO
      END DO
      !$OMP END DO
      ! update q and, if update_mass, update mass
      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         ! this VLOOP iterates over primal cell edges
         SELECT CASE (primal_deg(ij))
         CASE (4)
            edge1 = primal_edge(1, ij)
            edge2 = primal_edge(2, ij)
            edge3 = primal_edge(3, ij)
            edge4 = primal_edge(4, ij)
            sign1 = primal_ne(1, ij)
            sign2 = primal_ne(2, ij)
            sign3 = primal_ne(3, ij)
            sign4 = primal_ne(4, ij)
            !DIR$ SIMD
            DO l = 1, llm
               dmass = 0.
               dq = 0.
               dmass = dmass + sign1*fluxp(l, edge1)
               dq = dq + sign1*fluxq(l, edge1)
               dmass = dmass + sign2*fluxp(l, edge2)
               dq = dq + sign2*fluxq(l, edge2)
               dmass = dmass + sign3*fluxp(l, edge3)
               dq = dq + sign3*fluxq(l, edge3)
               dmass = dmass + sign4*fluxp(l, edge4)
               dq = dq + sign4*fluxq(l, edge4)
               newmass = massp(l, ij) - dmass/Ai(ij)
               massq(l, ij) = qi(l, ij)*massp(l, ij)
               qi(l, ij) = (massq(l, ij) - dq/Ai(ij))/newmass
               IF (update_mass) massp(l, ij) = newmass
            END DO
         CASE (5)
            edge1 = primal_edge(1, ij)
            edge2 = primal_edge(2, ij)
            edge3 = primal_edge(3, ij)
            edge4 = primal_edge(4, ij)
            edge5 = primal_edge(5, ij)
            sign1 = primal_ne(1, ij)
            sign2 = primal_ne(2, ij)
            sign3 = primal_ne(3, ij)
            sign4 = primal_ne(4, ij)
            sign5 = primal_ne(5, ij)
            !DIR$ SIMD
            DO l = 1, llm
               dmass = 0.
               dq = 0.
               dmass = dmass + sign1*fluxp(l, edge1)
               dq = dq + sign1*fluxq(l, edge1)
               dmass = dmass + sign2*fluxp(l, edge2)
               dq = dq + sign2*fluxq(l, edge2)
               dmass = dmass + sign3*fluxp(l, edge3)
               dq = dq + sign3*fluxq(l, edge3)
               dmass = dmass + sign4*fluxp(l, edge4)
               dq = dq + sign4*fluxq(l, edge4)
               dmass = dmass + sign5*fluxp(l, edge5)
               dq = dq + sign5*fluxq(l, edge5)
               newmass = massp(l, ij) - dmass/Ai(ij)
               massq(l, ij) = qi(l, ij)*massp(l, ij)
               qi(l, ij) = (massq(l, ij) - dq/Ai(ij))/newmass
               IF (update_mass) massp(l, ij) = newmass
            END DO
         CASE (6)
            edge1 = primal_edge(1, ij)
            edge2 = primal_edge(2, ij)
            edge3 = primal_edge(3, ij)
            edge4 = primal_edge(4, ij)
            edge5 = primal_edge(5, ij)
            edge6 = primal_edge(6, ij)
            sign1 = primal_ne(1, ij)
            sign2 = primal_ne(2, ij)
            sign3 = primal_ne(3, ij)
            sign4 = primal_ne(4, ij)
            sign5 = primal_ne(5, ij)
            sign6 = primal_ne(6, ij)
            !DIR$ SIMD
            DO l = 1, llm
               dmass = 0.
               dq = 0.
               dmass = dmass + sign1*fluxp(l, edge1)
               dq = dq + sign1*fluxq(l, edge1)
               dmass = dmass + sign2*fluxp(l, edge2)
               dq = dq + sign2*fluxq(l, edge2)
               dmass = dmass + sign3*fluxp(l, edge3)
               dq = dq + sign3*fluxq(l, edge3)
               dmass = dmass + sign4*fluxp(l, edge4)
               dq = dq + sign4*fluxq(l, edge4)
               dmass = dmass + sign5*fluxp(l, edge5)
               dq = dq + sign5*fluxq(l, edge5)
               dmass = dmass + sign6*fluxp(l, edge6)
               dq = dq + sign6*fluxq(l, edge6)
               newmass = massp(l, ij) - dmass/Ai(ij)
               massq(l, ij) = qi(l, ij)*massp(l, ij)
               qi(l, ij) = (massq(l, ij) - dq/Ai(ij))/newmass
               IF (update_mass) massp(l, ij) = newmass
            END DO
         CASE DEFAULT
            !DIR$ SIMD
            DO l = 1, llm
               dmass = 0.
               dq = 0.
               DO iedge = 1, primal_deg(ij)
                  edge = primal_edge(iedge, ij)
                  dmass = dmass + primal_ne(iedge, ij)*fluxp(l, edge)
                  dq = dq + primal_ne(iedge, ij)*fluxq(l, edge)
               END DO
               newmass = massp(l, ij) - dmass/Ai(ij)
               massq(l, ij) = qi(l, ij)*massp(l, ij)
               qi(l, ij) = (massq(l, ij) - dq/Ai(ij))/newmass
               IF (update_mass) massp(l, ij) = newmass
            END DO
         END SELECT
      END DO
      !$OMP END DO
      !---------------------------- advect_horiz ----------------------------------
      !--------------------------------------------------------------------------

      CALL trace_end("compute_advect_horiz")

   END SUBROUTINE compute_advect_horiz_unst

   ! Horizontal transport (S. Dubey, T. Dubos, Y. Meurdesoif)
   ! Slope-limited van Leer approach with hexagons
   SUBROUTINE compute_advect_horiz_hex(update_mass, diagflux_on, massp, fluxp, massq, fluxq, &
                                       qi, gradq3d, qfluxt, cc, Ai, xyz_i, qsat)
      USE trace
      USE omp_para
      USE abort_mod
      USE dimensions
      USE geometry, hide_Ai => Ai, hide_xyz_i => xyz_i
      USE grid_param, ONLY: llm, nqtot

      LOGICAL, INTENT(IN)       :: update_mass, diagflux_on
      REAL(rstd), INTENT(IN)    :: fluxp(3*iim*jjm, llm)      ! parent mass flux
      REAL(rstd), INTENT(OUT)   :: fluxq(3*iim*jjm, llm)      ! tracer mass flux
      REAL(rstd), INTENT(INOUT) :: massp(iim*jjm, llm), massq(iim*jjm, llm) ! parent and tracer mass
      REAL(rstd), INTENT(INOUT) :: qi(iim*jjm, llm)           ! tracer concentration
      REAL(rstd), INTENT(IN)    :: gradq3d(iim*jjm, llm, 3)   ! tracer gradient
      REAL(rstd), INTENT(INOUT) :: qfluxt(3*iim*jjm, llm)     ! time-integrated tracer flux
      REAL(rstd), INTENT(IN)    :: cc(3*iim*jjm, llm, 3)      ! barycenter of quadrilateral, where q is evaluated (1-point quadrature)
      REAL(rstd), INTENT(IN)    :: Ai(iim*jjm)
      REAL(rstd), INTENT(IN)    :: xyz_i(iim*jjm, 3)
      REAL(rstd), INTENT(IN), OPTIONAL    :: qsat(iim*jjm, llm)  ! saturating water vapor in kg/kg, used only for vanleer_qsat scheme

      REAL(rstd) :: dq, dmass, qe, newmass
      INTEGER :: ij, l, ij_tmp
      LOGICAL :: ok_qsat
      
      ok_qsat=PRESENT(qsat)

      IF (diagflux_on) CALL abort_acc("compute_advect_horiz : diagflux_on")

      CALL trace_start("compute_advect_horiz")
!---------------------------- advect_horiz ----------------------------------
!$acc data present(fluxp(:,:), fluxq(:,:), massp(:,:), massq(:,:), qi(:,:), gradq3d(:,:,:), cc(:,:,:), Ai(:), xyz_i(:,:), qsat(:,:)) async

      ! evaluate tracer field at point cc using piecewise linear reconstruction
      ! q(cc)= q0 + gradq.(cc-xyz_i) with xi centroid of hexagon
      ! sign*hqflux>0 iff outgoing
      !$acc parallel loop collapse(2) async
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            IF (ne_right*fluxp(ij + u_right, l) > 0.) THEN
               ij_tmp = ij
            ELSE
               ij_tmp = ij + t_right
            END IF

            qe = qi(ij_tmp, l)
            qe = qe + (cc(ij + u_right, l, 1) - xyz_i(ij_tmp, 1))*gradq3d(ij_tmp, l, 1)
            qe = qe + (cc(ij + u_right, l, 2) - xyz_i(ij_tmp, 2))*gradq3d(ij_tmp, l, 2)
            qe = qe + (cc(ij + u_right, l, 3) - xyz_i(ij_tmp, 3))*gradq3d(ij_tmp, l, 3)
            
            IF (ok_qsat) qe = MIN(qsat(ij_tmp,l), qe)
            
            fluxq(ij + u_right, l) = fluxp(ij + u_right, l)*qe

            IF (ne_lup*fluxp(ij + u_lup, l) > 0.) THEN
               ij_tmp = ij
            ELSE
               ij_tmp = ij + t_lup
            END IF

            qe = qi(ij_tmp, l)
            qe = qe + (cc(ij + u_lup, l, 1) - xyz_i(ij_tmp, 1))*gradq3d(ij_tmp, l, 1)
            qe = qe + (cc(ij + u_lup, l, 2) - xyz_i(ij_tmp, 2))*gradq3d(ij_tmp, l, 2)
            qe = qe + (cc(ij + u_lup, l, 3) - xyz_i(ij_tmp, 3))*gradq3d(ij_tmp, l, 3)
            
            IF (ok_qsat) qe = MIN(qsat(ij_tmp,l), qe)

            fluxq(ij + u_lup, l) = fluxp(ij + u_lup, l)*qe

            IF (ne_ldown*fluxp(ij + u_ldown, l) > 0.) THEN
               ij_tmp = ij
            ELSE
               ij_tmp = ij + t_ldown
            END IF

            qe = qi(ij_tmp, l)

            qe = qe + (cc(ij + u_ldown, l, 1) - xyz_i(ij_tmp, 1))*gradq3d(ij_tmp, l, 1)
            qe = qe + (cc(ij + u_ldown, l, 2) - xyz_i(ij_tmp, 2))*gradq3d(ij_tmp, l, 2)
            qe = qe + (cc(ij + u_ldown, l, 3) - xyz_i(ij_tmp, 3))*gradq3d(ij_tmp, l, 3)
            IF (ok_qsat) qe = MIN(qsat(ij_tmp,l), qe)
            fluxq(ij + u_ldown, l) = fluxp(ij + u_ldown, l)*qe
         END DO
      END DO

      IF (diagflux_on) THEN
         !$acc parallel loop collapse(2) copy(qfluxt(:,:)) async
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin_ext, ij_end_ext
               qfluxt(ij + u_right, l) = qfluxt(ij + u_right, l) + fluxq(ij + u_right, l)
               qfluxt(ij + u_lup, l) = qfluxt(ij + u_lup, l) + fluxq(ij + u_lup, l)
               qfluxt(ij + u_ldown, l) = qfluxt(ij + u_ldown, l) + fluxq(ij + u_ldown, l)
            END DO
         END DO
      END IF

      ! update q and, if update_mass, update mass
      !$acc parallel loop collapse(2) async
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin, ij_end
            dmass = 0.
            dq = 0.
            dmass = dmass + ne_rup*fluxp(ij + u_rup, l)
            dq = dq + ne_rup*fluxq(ij + u_rup, l)
            dmass = dmass + ne_lup*fluxp(ij + u_lup, l)
            dq = dq + ne_lup*fluxq(ij + u_lup, l)
            dmass = dmass + ne_left*fluxp(ij + u_left, l)
            dq = dq + ne_left*fluxq(ij + u_left, l)
            dmass = dmass + ne_ldown*fluxp(ij + u_ldown, l)
            dq = dq + ne_ldown*fluxq(ij + u_ldown, l)
            dmass = dmass + ne_rdown*fluxp(ij + u_rdown, l)
            dq = dq + ne_rdown*fluxq(ij + u_rdown, l)
            dmass = dmass + ne_right*fluxp(ij + u_right, l)
            dq = dq + ne_right*fluxq(ij + u_right, l)
            newmass = massp(ij, l) - dmass/Ai(ij)
            massq(ij, l) = qi(ij, l)*massp(ij, l)
            qi(ij, l) = (massq(ij, l) - dq/Ai(ij))/newmass
            IF (update_mass) massp(ij, l) = newmass
         END DO
      END DO
!$acc end data
!---------------------------- advect_horiz ----------------------------------
      CALL trace_end("compute_advect_horiz")

   END SUBROUTINE compute_advect_horiz_hex

END SUBMODULE compute_advect_horiz_mod
