MODULE euler_scheme_mod
   IMPLICIT NONE
   PRIVATE

   PUBLIC :: euler_scheme

CONTAINS

   SUBROUTINE Euler_scheme(with_dps, fluxt_zero)
      USE icosa
      USE disvert_mod
      USE omp_para
      USE trace
      USE abort_mod, ONLY: abort_acc
      USE time_scheme_base_mod, ONLY: &
         compute_accumulate_fluxes, &
         f_ps, f_dps, f_mass, f_dmass, &
         f_hfluxt, f_wfluxt, f_hflux, f_wflux, &
         f_u, f_theta_rhodz, f_du, f_dtheta_rhodz

      LOGICAL :: with_dps
      LOGICAL, OPTIONAL :: fluxt_zero(ndomain) ! set to .TRUE. to start accumulating fluxes in time
      REAL(rstd), POINTER :: ps(:), dps(:)
      REAL(rstd), POINTER :: u(:, :), du(:, :)
      REAL(rstd), POINTER :: mass(:, :), dmass(:, :)
      REAL(rstd), POINTER :: theta_rhodz(:, :, :), dtheta_rhodz(:, :, :)
      REAL(rstd), POINTER :: hflux(:, :), wflux(:, :), hfluxt(:, :), wfluxt(:, :)
      INTEGER :: ind
      INTEGER :: ij, l
      CALL trace_start("Euler_scheme")

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)

         IF (with_dps) THEN ! update ps/mass
            CALL abort_acc("Euler_scheme/with_dps")
            IF (caldyn_eta == eta_mass) THEN ! update ps
               ps = f_ps(ind); dps = f_dps(ind)
               IF (is_omp_first_level) THEN
                  !DIR$ SIMD
                  DO ij = ij_begin, ij_end
                     ps(ij) = ps(ij) + dt*dps(ij)
                  END DO
               END IF
            ELSE ! update mass
               mass = f_mass(ind); dmass = f_dmass(ind)
               DO l = ll_begin, ll_end
                  !DIR$ SIMD
                  DO ij = ij_begin, ij_end
                     mass(ij, l) = mass(ij, l) + dt*dmass(ij, l)
                  END DO
               END DO
            END IF

            hflux = f_hflux(ind); hfluxt = f_hfluxt(ind)
            wflux = f_wflux(ind); wfluxt = f_wfluxt(ind)
            CALL compute_accumulate_fluxes(hflux, wflux, hfluxt, wfluxt, dt, fluxt_zero(ind))
         END IF ! update ps/mass

         u = f_u(ind); theta_rhodz = f_theta_rhodz(ind)
         du = f_du(ind); dtheta_rhodz = f_dtheta_rhodz(ind)

         CALL compute_euler_scheme(u, du, theta_rhodz, dtheta_rhodz)
      END DO

      CALL trace_end("Euler_scheme")

   CONTAINS

      SUBROUTINE compute_euler_scheme(u, du, theta_rhodz, dtheta_rhodz)
         REAL(rstd), INTENT(INOUT) :: u(iim*3*jjm, llm)
         REAL(rstd), INTENT(IN)    :: du(iim*3*jjm, llm)
         REAL(rstd), INTENT(INOUT) :: theta_rhodz(iim*jjm, llm, nqdyn)
         REAL(rstd), INTENT(IN)    :: dtheta_rhodz(iim*jjm, llm, nqdyn)

         !$acc data present(theta_rhodz(:,:,:), u(:,:),du(:,:), dtheta_rhodz(:,:,:)) async

         !$acc parallel loop async
         DO l = ll_begin, ll_end
            !$acc loop
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               u(ij + u_right, l) = u(ij + u_right, l) + dt*du(ij + u_right, l)
               u(ij + u_lup, l) = u(ij + u_lup, l) + dt*du(ij + u_lup, l)
               u(ij + u_ldown, l) = u(ij + u_ldown, l) + dt*du(ij + u_ldown, l)
               theta_rhodz(ij, l, 1) = theta_rhodz(ij, l, 1) + dt*dtheta_rhodz(ij, l, 1)
            END DO
         END DO

         !$acc end data
      END SUBROUTINE compute_euler_scheme

   END SUBROUTINE Euler_scheme

END MODULE euler_scheme_mod
