MODULE compute_velocity_mod
  USE icosa
  USE caldyn_vars_mod
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: velocity

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
    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
  
  SUBROUTINE compute_velocity(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

END MODULE compute_velocity_mod