MODULE compute_velocity_mod
   USE icosa
   USE caldyn_vars_mod
   IMPLICIT NONE
   PRIVATE

   PUBLIC :: velocity, compute_velocity_hex, compute_velocity_unst

CONTAINS

   !------------------- Conversion from momentum to horizontal velocity  ------------------

   SUBROUTINE velocity(f_geopot, f_ps, f_rhodz, f_u, f_W, f_buf_Fel, f_uh, f_uz)
      USE disvert_mod, ONLY: caldyn_eta, eta_mass
      USE compute_diagnostics_mod, ONLY: compute_rhodz, compute_velocity
      TYPE(t_field), POINTER :: f_geopot(:), f_ps(:), f_rhodz(:), &
                                f_u(:), f_W(:), &  ! IN
                                f_buf_Fel(:), &          ! BUF
                                f_uh(:), f_uz(:)      ! OUT
      REAL(rstd), POINTER :: geopot(:, :), ps(:), rhodz(:, :), u(:, :), W(:, :), uh(:, :), uz(:, :), F_el(:, :)
      INTEGER :: ind

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         geopot = f_geopot(ind)
         rhodz = f_rhodz(ind)
         u = f_u(ind)
         W = f_W(ind)
         uh = f_uh(ind)
         F_el = f_buf_Fel(ind)
         IF (caldyn_eta == eta_mass) THEN
            ps = f_ps(ind)
            CALL compute_rhodz(.TRUE., ps, rhodz)
         END IF
         uz = f_uz(ind)
         !$OMP BARRIER
         CALL compute_velocity(geopot, rhodz, u, W, F_el, uh, uz)
         !$OMP BARRIER
      END DO
   END SUBROUTINE velocity

!-------------- Wrappers for F2008 conformity -----------------

   SUBROUTINE compute_velocity_unst(Phi, rhodz, u, W, F_el, uh, uz)
      REAL(rstd), INTENT(IN)  :: Phi(:, :), u(:, :), rhodz(:, :), W(:, :)
      REAL(rstd), INTENT(OUT) :: F_el(:, :), uh(:, :), uz(:, :)
      CALL compute_velocity_unst_(Phi, rhodz, u, W, F_el, uh, uz)
   END SUBROUTINE compute_velocity_unst

   SUBROUTINE compute_velocity_hex(Phi, rhodz, u, W, F_el, uh, uz)
      REAL(rstd), INTENT(IN)  :: Phi(:, :), u(:, :), rhodz(:, :), W(:, :)
      REAL(rstd), INTENT(OUT) :: F_el(:, :), uh(:, :), uz(:, :)
      CALL compute_velocity_hex_(Phi, rhodz, u, W, F_el, uh, uz)
   END SUBROUTINE compute_velocity_hex

!--------------------------------------------------------------

   SUBROUTINE compute_velocity_unst_(Phi, rhodz, u, W, F_el, uh, uz)
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num)    :: Phi, W, uz
      REAL(C_DOUBLE), DIMENSION(llm, primal_num) :: rhodz
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)    :: u, uh
      REAL(C_DOUBLE), DIMENSION(llm + 1, edge_num)   :: F_el
      IF (hydrostatic) THEN
         uh(:, :) = u(:, :)
         uz(:, :) = 0.
      ELSE
         STOP 'Non-hydrostatic compute_velocity not implemented yet'
      END IF
   END SUBROUTINE compute_velocity_unst_

   SUBROUTINE compute_velocity_hex_(Phi, rhodz, u, W, F_el, uh, uz)
      USE omp_para
      REAL(rstd), INTENT(IN) :: Phi(iim*jjm, llm + 1)
      REAL(rstd), INTENT(IN) :: rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(IN) :: u(3*iim*jjm, llm)
      REAL(rstd), INTENT(IN) :: W(iim*jjm, llm + 1)
      REAL(rstd), INTENT(OUT) :: uh(3*iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: uz(iim*jjm, llm)
      INTEGER :: ij, l
      REAL(rstd) :: F_el(3*iim*jjm, llm + 1)
      REAL(rstd) :: uu_right, uu_lup, uu_ldown, W_el, DePhil
      ! NB : u and uh are not in DEC form, they are normal components
      ! => we must divide by de
      IF (hydrostatic) THEN
         uh(:, :) = u(:, :)
         uz(:, :) = 0.
      ELSE
         DO l = ll_begin, ll_endp1 ! compute on l levels (interfaces)
            DO ij = ij_begin_ext, ij_end_ext
               ! Compute on edge 'right'
               W_el = .5*(W(ij, l) + W(ij + t_right, l))
               DePhil = ne_right*(Phi(ij + t_right, l) - Phi(ij, l))
               F_el(ij + u_right, l) = DePhil*W_el/de(ij + u_right)
               ! Compute on edge 'lup'
               W_el = .5*(W(ij, l) + W(ij + t_lup, l))
               DePhil = ne_lup*(Phi(ij + t_lup, l) - Phi(ij, l))
               F_el(ij + u_lup, l) = DePhil*W_el/de(ij + u_lup)
               ! Compute on edge 'ldown'
               W_el = .5*(W(ij, l) + W(ij + t_ldown, l))
               DePhil = ne_ldown*(Phi(ij + t_ldown, l) - Phi(ij, l))
               F_el(ij + u_ldown, l) = DePhil*W_el/de(ij + u_ldown)
            END DO
         END DO

         ! We need a barrier here because we compute F_el above and do a vertical average below
         !$OMP BARRIER

         DO l = ll_begin, ll_end ! compute on k levels (full levels)
            DO ij = ij_begin_ext, ij_end_ext
               ! w = vertical momentum = g^-2*dPhi/dt = uz/g
               uz(ij, l) = (.5*g)*(W(ij, l) + W(ij, l + 1))/rhodz(ij, l)
               ! uh = u-w.grad(Phi) = u - uz.grad(z)
               uh(ij + u_right, l) = u(ij + u_right, l) - &
                                     (F_el(ij + u_right, l) + F_el(ij + u_right, l + 1))/(rhodz(ij, l) + rhodz(ij + t_right, l))
               uh(ij + u_lup, l) = u(ij + u_lup, l) - &
                                   (F_el(ij + u_lup, l) + F_el(ij + u_lup, l + 1))/(rhodz(ij, l) + rhodz(ij + t_lup, l))
               uh(ij + u_ldown, l) = u(ij + u_ldown, l) - &
                                     (F_el(ij + u_ldown, l) + F_el(ij + u_ldown, l + 1))/(rhodz(ij, l) + rhodz(ij + t_ldown, l))
            END DO
         END DO

      END IF
   END SUBROUTINE compute_velocity_hex_

END MODULE compute_velocity_mod
