MODULE caldyn_hevi_mod
   IMPLICIT NONE
   PRIVATE
   PUBLIC caldyn_hevi

CONTAINS

   SUBROUTINE caldyn_hevi(write_out, tau, f_phis, f_ps, f_mass, f_theta_rhodz, f_u, f_q, &
      f_W, f_geopot, f_hflux, f_wflux, f_dps, f_dmass, f_dtheta_rhodz, &
      f_du_slow, f_du_fast, f_dPhi_slow, f_dPhi_fast, f_dW_slow, f_dW_fast)
      USE icosa
      USE abort_mod, ONLY: abort_acc
      USE trace

      USE compute_caldyn_mod, ONLY : compute_pvort_only, compute_theta, &
         compute_geopot, compute_caldyn_fast, compute_caldyn_slow_hydro, &
         compute_caldyn_coriolis, compute_caldyn_vert, compute_caldyn_Kv, &
         compute_caldyn_vert_NH, compute_caldyn_slow_NH, compute_caldyn_solver

      USE observable_mod, ONLY: f_theta
      USE disvert_mod, ONLY : caldyn_eta, eta_mass

      USE caldyn_vars_mod, ONLY: hydrostatic
      USE caldyn_vars_mod, ONLY: f_qu, f_qv, f_pk, f_wwuu, f_berni, f_Ftheta, f_Fel, f_gradPhi2, f_wil, f_Wetadot
      USE caldyn_vars_mod, ONLY: req_ps, req_mass, req_theta_rhodz, req_u, req_qu, req_geopot, req_w

      USE transfert_mod

      LOGICAL, INTENT(IN)    :: write_out
      REAL(rstd), INTENT(IN) :: tau
      TYPE(t_field), POINTER :: f_phis(:)
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_mass(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_u(:)
      TYPE(t_field), POINTER :: f_q(:)
      TYPE(t_field), POINTER :: f_W(:)
      TYPE(t_field), POINTER :: f_geopot(:)
      TYPE(t_field), POINTER :: f_hflux(:), f_wflux(:)
      TYPE(t_field) :: f_dps(:)
      TYPE(t_field) :: f_dmass(:)
      TYPE(t_field) :: f_dtheta_rhodz(:)
      TYPE(t_field) :: f_du_slow(:)
      TYPE(t_field) :: f_du_fast(:)
      TYPE(t_field) :: f_dW_slow(:)
      TYPE(t_field) :: f_dW_fast(:)
      TYPE(t_field) :: f_dPhi_slow(:)
      TYPE(t_field) :: f_dPhi_fast(:)

      REAL(rstd), POINTER :: ps(:), dps(:), phis(:)
      REAL(rstd), POINTER :: mass(:, :), theta_rhodz(:, :, :), dtheta_rhodz(:, :, :)
      REAL(rstd), POINTER :: du(:, :), dW(:, :), dPhi(:, :), hflux(:, :), wflux(:, :)
      REAL(rstd), POINTER :: u(:, :), w(:, :), qu(:, :), qv(:, :)
      REAL(rstd), POINTER :: berni(:,:), Ftheta(:,:), Kv(:, :), hv(:, :)
! temporary shared variable
      REAL(rstd), POINTER :: theta(:, :, :)
      REAL(rstd), POINTER :: pk(:, :, :)
      REAL(rstd), POINTER :: geopot(:, :)
      REAL(rstd), POINTER :: convm(:, :)
      REAL(rstd), POINTER :: wwuu(:, :)
      ! Non-hydrostatic, not on GPU
      REAL(rstd), POINTER :: F_el(:, :), gradPhi2(:, :), w_il(:, :), W_etadot(:, :), pres(:, :), m_il(:, :)
      INTEGER :: ind
      LOGICAL, SAVE :: first = .TRUE.
!$OMP THREADPRIVATE(first)

      IF (first) THEN
         first = .FALSE.
         IF (caldyn_eta == eta_mass) THEN
            CALL init_message(f_ps, req_i1, req_ps)
         ELSE
            CALL init_message(f_mass, req_i1, req_mass)
         END IF
         CALL init_message(f_theta_rhodz, req_i1, req_theta_rhodz)
         CALL init_message(f_u, req_e1_vect, req_u)
         CALL init_message(f_qu, req_e1_scal, req_qu)
!         IF (caldyn_kinetic == kinetic_consistent) CALL init_message(f_Kv,req_z1_scal,req_Kv)
         IF (.NOT. hydrostatic) THEN
            CALL init_message(f_geopot, req_i1, req_geopot)
            CALL init_message(f_w, req_i1, req_w)
         END IF
      END IF

      CALL trace_start("caldyn")

      IF (caldyn_eta == eta_mass) THEN
         CALL send_message(f_ps, req_ps) ! COM00
         CALL wait_message(req_ps) ! COM00
      ELSE
         CALL abort_acc("HEVI_scheme/!eta_mass")
         CALL send_message(f_mass, req_mass) ! COM00
         CALL wait_message(req_mass) ! COM00
      END IF
      CALL send_message(f_theta_rhodz, req_theta_rhodz) ! COM01
      CALL wait_message(req_theta_rhodz) ! COM01 Moved from caldyn_pvort

      IF (.NOT. hydrostatic) THEN
         CALL abort_acc("HEVI_scheme/!hydrostatic")
         CALL send_message(f_geopot, req_geopot) ! COM03
         CALL wait_message(req_geopot) ! COM03
         CALL send_message(f_w, req_w) ! COM04
         CALL wait_message(req_w) ! COM04
      END IF

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         theta_rhodz = f_theta_rhodz(ind)
         mass = f_mass(ind)
         theta = f_theta(ind)
         CALL compute_theta(ps, theta_rhodz, mass, theta)
         pk = f_pk(ind)
         geopot = f_geopot(ind)
         du = f_du_fast(ind)
         berni = f_berni(ind)

         IF (hydrostatic) THEN
            !$acc kernels present(du) async
            du(:, :) = 0.0d0
            !$acc end kernels

            CALL compute_geopot(mass, theta, ps, berni, pk, geopot)
         ELSE
            CALL abort_acc("HEVI_scheme/!hydrostatic")
            phis = f_phis(ind)
            W = f_W(ind)
            dW = f_dW_fast(ind)
            dPhi = f_dPhi_fast(ind)
            ! reuse buffers
            m_il = f_wil(ind)
            pres = f_gradPhi2(ind)
            CALL compute_caldyn_solver(tau, phis, mass, theta, pk(:,:,1), geopot, W, m_il, pres, dPhi, dW, du) ! computes d(Phi,W,du)_fast and updates Phi,W
         END IF
         u = f_u(ind)
         CALL compute_caldyn_fast(tau, theta, geopot, pk, berni, du, u) ! computes du_fast and updates u
      END DO

      CALL send_message(f_u, req_u) ! COM02
      CALL wait_message(req_u)   ! COM02

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         u = f_u(ind)
         mass = f_mass(ind)
         qu = f_qu(ind)
         qv = f_qv(ind)
!         hv = f_hv(ind)
!       PRINT *, 'caldyn_hevi : u, mass, qu qv, hv', &
!            SHAPE(u), SHAPE(mass), SHAPE(qu), SHAPE(qv), SHAPE(hv)
         CALL compute_pvort_only(u, mass, qu, qv)
!         IF (caldyn_kinetic == kinetic_consistent) THEN
!            Kv = f_Kv(ind)
!            CALL compute_caldyn_Kv(u, Kv)
!         END IF
      END DO

      CALL send_message(f_qu, req_qu) ! COM03
      CALL wait_message(req_qu) ! COM03

!      IF(caldyn_kinetic==kinetic_consistent) THEN
!         CALL transfert_request(f_hv,req_z1_scal)
!         CALL send_message(f_Kv,req_Kv)
!         CALL wait_message(req_Kv)
!      END IF

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         u = f_u(ind)
         mass = f_mass(ind)
         theta = f_theta(ind)
         qu = f_qu(ind)
         hflux = f_hflux(ind)
         convm = f_dmass(ind)
         dtheta_rhodz = f_dtheta_rhodz(ind)
         du = f_du_slow(ind)
         berni = f_berni(ind) ! buffer for partial Bernoulli function
         Ftheta = f_Ftheta(ind)

         IF (hydrostatic) THEN
            CALL compute_caldyn_slow_hydro(.TRUE., u, mass, berni, hflux, du)
         ELSE
            CALL abort_acc("HEVI_scheme/!hydrostatic")
            W = f_W(ind)
            dW = f_dW_slow(ind)
            geopot = f_geopot(ind)
            dPhi = f_dPhi_slow(ind)
            F_el = f_Fel(ind)
            gradPhi2 = f_gradPhi2(ind)
            w_il = f_wil(ind)
            CALL compute_caldyn_slow_NH(u, mass, geopot, W, F_el, gradPhi2, w_il, hflux, du, dPhi, dW)
         END IF
         CALL compute_caldyn_Coriolis(hflux, theta, qu, Ftheta, convm, dtheta_rhodz, du)
         IF (caldyn_eta == eta_mass) THEN
            wflux = f_wflux(ind)
            wwuu = f_wwuu(ind)
            dps = f_dps(ind)
            CALL compute_caldyn_vert(u, theta, mass, convm, wflux, wwuu, dps, dtheta_rhodz, du)
            IF (.NOT. hydrostatic) THEN
               CALL abort_acc("HEVI_scheme/!hydrostatic")
               W_etadot = f_Wetadot(ind)
               CALL compute_caldyn_vert_NH(mass, geopot, W, wflux, W_etadot, du, dPhi, dW)
            END IF
         END IF
      END DO

!$OMP BARRIER
      !    CALL check_mass_conservation(f_ps,f_dps)
      CALL trace_end("caldyn_hevi")
!!$OMP BARRIER

   END SUBROUTINE caldyn_hevi

END MODULE caldyn_hevi_mod
