MODULE compute_caldyn_slow_NH_mod
  USE grid_param, ONLY : llm
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: compute_caldyn_slow_NH

CONTAINS

  SUBROUTINE compute_caldyn_slow_NH(u,rhodz,Phi,W, F_el,gradPhi2,w_il, hflux,du,dPhi,dW)
    USE icosa
    USE trace
    USE caldyn_vars_mod
    USE omp_para, ONLY : ll_begin, ll_end,ll_beginp1,ll_endp1
    REAL(rstd),INTENT(IN)  :: u(3*iim*jjm,llm)    ! prognostic "velocity"
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)  ! rho*dz
    REAL(rstd),INTENT(IN)  :: Phi(iim*jjm,llm+1)  ! prognostic geopotential
    REAL(rstd),INTENT(IN)  :: W(iim*jjm,llm+1)    ! prognostic vertical momentum

    REAL(rstd),INTENT(OUT) :: hflux(3*iim*jjm,llm) ! hflux in kg/s
    REAL(rstd),INTENT(OUT) :: du(3*iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: dW(iim*jjm,llm+1)
    REAL(rstd),INTENT(OUT) :: dPhi(iim*jjm,llm+1)

    REAL(rstd) :: w_il(iim*jjm,llm+1) ! Wil/mil
    REAL(rstd) :: F_el(3*iim*jjm,llm+1) ! NH mass flux
    REAL(rstd) :: gradPhi2(iim*jjm,llm+1) ! grad_Phi**2
    REAL(rstd) :: DePhil(3*iim*jjm,llm+1) ! grad(Phi)
    
    INTEGER :: ij,l,kdown,kup
    REAL(rstd) :: W_el, W2_el, uu_right, uu_lup, uu_ldown, gPhi2, dP, divG, u2, uu

    REAL(rstd) :: berni(iim*jjm,llm)  ! Bernoulli function
    REAL(rstd) :: G_el(3*iim*jjm,llm+1) ! horizontal flux of W
    REAL(rstd) :: v_el(3*iim*jjm,llm+1)

    REAL(rstd) :: berni1(iim*jjm)  ! Bernoulli function
    REAL(rstd) :: G_el1(3*iim*jjm) ! horizontal flux of W
    REAL(rstd) :: v_el1(3*iim*jjm)

    CALL trace_start("compute_caldyn_slow_NH")

    IF(dysl) THEN

!$OMP BARRIER
#include "../kernels_hex/caldyn_slow_NH.k90"
!$OMP BARRIER
     
     ELSE

#define BERNI(ij) berni1(ij)
#define G_EL(ij) G_el1(ij)
#define V_EL(ij) v_el1(ij)

    DO l=ll_begin, ll_endp1 ! compute on l levels (interfaces)
       IF(l==1) THEN
          kdown=1
       ELSE
          kdown=l-1
       END IF
       IF(l==llm+1) THEN
          kup=llm
       ELSE
          kup=l
       END IF
       ! below : "checked" means "formula also valid when kup=kdown (top/bottom)"
       ! compute mil, wil=Wil/mil
       DO ij=ij_begin_ext, ij_end_ext
          w_il(ij,l) = 2.*W(ij,l)/(rhodz(ij,kdown)+rhodz(ij,kup)) ! checked
       END DO
       ! compute DePhi, v_el, G_el, F_el
       ! v_el, W2_el and therefore G_el incorporate metric factor le_de
       ! while DePhil, W_el and F_el don't
       DO ij=ij_begin_ext, ij_end_ext
          ! Compute on edge 'right'
          W_el  = .5*( W(ij,l)+W(ij+t_right,l) )
          DePhil(ij+u_right,l) = ne_right*(Phi(ij+t_right,l)-Phi(ij,l))
          F_el(ij+u_right,l)   = DePhil(ij+u_right,l)*W_el
          W2_el = .5*le_de(ij+u_right) * &
               ( W(ij,l)*w_il(ij,l) + W(ij+t_right,l)*w_il(ij+t_right,l) )
          V_EL(ij+u_right) = .5*le_de(ij+u_right)*(u(ij+u_right,kup)+u(ij+u_right,kdown)) ! checked
          G_EL(ij+u_right) = V_EL(ij+u_right)*W_el - DePhil(ij+u_right,l)*W2_el
          ! Compute on edge 'lup'
          W_el  = .5*( W(ij,l)+W(ij+t_lup,l) )
          DePhil(ij+u_lup,l) = ne_lup*(Phi(ij+t_lup,l)-Phi(ij,l))
          F_el(ij+u_lup,l)   = DePhil(ij+u_lup,l)*W_el
          W2_el = .5*le_de(ij+u_lup) * &
               ( W(ij,l)*w_il(ij,l) + W(ij+t_lup,l)*w_il(ij+t_lup,l) )
          V_EL(ij+u_lup) = .5*le_de(ij+u_lup)*( u(ij+u_lup,kup) + u(ij+u_lup,kdown)) ! checked
          G_EL(ij+u_lup) = V_EL(ij+u_lup)*W_el - DePhil(ij+u_lup,l)*W2_el
          ! Compute on edge 'ldown'
          W_el  = .5*( W(ij,l)+W(ij+t_ldown,l) )
          DePhil(ij+u_ldown,l) = ne_ldown*(Phi(ij+t_ldown,l)-Phi(ij,l))
          F_el(ij+u_ldown,l) = DePhil(ij+u_ldown,l)*W_el
          W2_el = .5*le_de(ij+u_ldown) * &
               ( W(ij,l)*w_il(ij,l) + W(ij+t_ldown,l)*w_il(ij+t_ldown,l) )
          V_EL(ij+u_ldown) = .5*le_de(ij+u_ldown)*( u(ij+u_ldown,kup) + u(ij+u_ldown,kdown)) ! checked
          G_EL(ij+u_ldown) = V_EL(ij+u_ldown)*W_el - DePhil(ij+u_ldown,l)*W2_el
       END DO
       ! compute GradPhi2, dPhi, dW
       DO ij=ij_begin_ext, ij_end_ext
          gradPhi2(ij,l) = &
               1/(2*Ai(ij))*(le_de(ij+u_right)*DePhil(ij+u_right,l)**2 + &
               le_de(ij+u_rup)*DePhil(ij+u_rup,l)**2 +      &
               le_de(ij+u_lup)*DePhil(ij+u_lup,l)**2 +      &
               le_de(ij+u_left)*DePhil(ij+u_left,l)**2 +    &
               le_de(ij+u_ldown)*DePhil(ij+u_ldown,l)**2 +  &
               le_de(ij+u_rdown)*DePhil(ij+u_rdown,l)**2 )

          dPhi(ij,l) = gradPhi2(ij,l)*w_il(ij,l) -1/(2*Ai(ij))* & 
               ( DePhil(ij+u_right,l)*V_EL(ij+u_right) + & ! -v.gradPhi, 
                 DePhil(ij+u_rup,l)*V_EL(ij+u_rup) +     & ! v_el already has le_de
                 DePhil(ij+u_lup,l)*V_EL(ij+u_lup) +     &
                 DePhil(ij+u_left,l)*V_EL(ij+u_left) +   &
                 DePhil(ij+u_ldown,l)*V_EL(ij+u_ldown) + &
                 DePhil(ij+u_rdown,l)*V_EL(ij+u_rdown) )

          dW(ij,l) = -1./Ai(ij)*(           & ! -div(G_el), 
               ne_right*G_EL(ij+u_right) +  & ! G_el already has le_de
               ne_rup*G_EL(ij+u_rup) +      &
               ne_lup*G_EL(ij+u_lup) +      &  
               ne_left*G_EL(ij+u_left) +    &
               ne_ldown*G_EL(ij+u_ldown) +  &
               ne_rdown*G_EL(ij+u_rdown))
       END DO
    END DO

    DO l=ll_begin, ll_end ! compute on k levels (layers)
       ! Compute berni at scalar points
       DO ij=ij_begin_ext, ij_end_ext
          BERNI(ij) = &
               1/(4*Ai(ij))*( &
                    le_de(ij+u_right)*u(ij+u_right,l)**2 +    &
                    le_de(ij+u_rup)*u(ij+u_rup,l)**2 +          &
                    le_de(ij+u_lup)*u(ij+u_lup,l)**2 +          &
                    le_de(ij+u_left)*u(ij+u_left,l)**2 +       &
                    le_de(ij+u_ldown)*u(ij+u_ldown,l)**2 +    &
                    le_de(ij+u_rdown)*u(ij+u_rdown,l)**2 ) &
               - .25*( gradPhi2(ij,l)  *w_il(ij,l)**2 +   &
                      gradPhi2(ij,l+1)*w_il(ij,l+1)**2 )
       END DO
       ! Compute mass flux and grad(berni) at edges
       DO ij=ij_begin_ext, ij_end_ext
          ! Compute on edge 'right'
          uu_right = 0.5*(rhodz(ij,l)+rhodz(ij+t_right,l))*u(ij+u_right,l) &
                    -0.5*(F_el(ij+u_right,l)+F_el(ij+u_right,l+1))
          hflux(ij+u_right,l) = uu_right*le_de(ij+u_right)
          du(ij+u_right,l) = ne_right*(BERNI(ij)-BERNI(ij+t_right))
          ! Compute on edge 'lup'
          uu_lup = 0.5*(rhodz(ij,l)+rhodz(ij+t_lup,l))*u(ij+u_lup,l) &
                  -0.5*(F_el(ij+u_lup,l)+F_el(ij+u_lup,l+1))
          hflux(ij+u_lup,l) = uu_lup*le_de(ij+u_lup)
          du(ij+u_lup,l) = ne_lup*(BERNI(ij)-BERNI(ij+t_lup))
          ! Compute on edge 'ldown'
          uu_ldown = 0.5*(rhodz(ij,l)+rhodz(ij+t_ldown,l))*u(ij+u_ldown,l) &
                    -0.5*(F_el(ij+u_ldown,l)+F_el(ij+u_ldown,l+1))
          hflux(ij+u_ldown,l) = uu_ldown*le_de(ij+u_ldown)
          du(ij+u_ldown,l) = ne_ldown*(BERNI(ij)-BERNI(ij+t_ldown))
       END DO
    END DO

#undef V_EL
#undef G_EL
#undef BERNI

    END IF ! dysl

    CALL trace_end("compute_caldyn_slow_NH")

  END SUBROUTINE compute_caldyn_slow_NH

END MODULE compute_caldyn_slow_NH_mod