MODULE compute_caldyn_vert_NH_mod
  USE prec, ONLY : rstd
  USE grid_param
  USE disvert_mod
  USE omp_para
  USE trace
  IMPLICIT NONE
  PRIVATE
  SAVE

  PUBLIC :: compute_caldyn_vert_nh_manual, &
       compute_caldyn_vert_nh_hex

CONTAINS

  SUBROUTINE compute_caldyn_vert_NH_hex(mass,geopot,W,wflux, W_etadot, du,dPhi,dW)
    REAL(rstd),INTENT(IN) :: mass(iim*jjm,llm)
    REAL(rstd),INTENT(IN) :: geopot(iim*jjm,llm+1)
    REAL(rstd),INTENT(IN) :: W(iim*jjm,llm+1)
    REAL(rstd),INTENT(IN) :: wflux(iim*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: dPhi(iim*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: dW(iim*jjm,llm+1)
    REAL(rstd) :: W_etadot(iim*jjm,llm) ! vertical flux of vertical momentum
    ! local arrays
    REAL(rstd) :: eta_dot(iim*jjm, llm) ! eta_dot in full layers
    REAL(rstd) :: wcov(iim*jjm,llm) ! covariant vertical momentum in full layers
    ! indices and temporary values
    INTEGER    :: ij, l
    REAL(rstd) :: wflux_ij, w_ij

    CALL trace_start("compute_caldyn_vert_nh")

!$OMP BARRIER
#include "../kernels_hex/caldyn_vert_NH.k90"
!$OMP BARRIER
    CALL trace_end("compute_caldyn_vert_nh")
  END SUBROUTINE compute_caldyn_vert_NH_hex

  SUBROUTINE compute_caldyn_vert_NH_manual(mass,geopot,W,wflux, W_etadot, du,dPhi,dW)
    REAL(rstd),INTENT(IN) :: mass(iim*jjm,llm)
    REAL(rstd),INTENT(IN) :: geopot(iim*jjm,llm+1)
    REAL(rstd),INTENT(IN) :: W(iim*jjm,llm+1)
    REAL(rstd),INTENT(IN) :: wflux(iim*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: dPhi(iim*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: dW(iim*jjm,llm+1)
    REAL(rstd) :: W_etadot(iim*jjm,llm) ! vertical flux of vertical momentum
    ! local arrays
    REAL(rstd) :: eta_dot(iim*jjm, llm) ! eta_dot in full layers
    REAL(rstd) :: wcov(iim*jjm,llm) ! covariant vertical momentum in full layers
    ! indices and temporary values
    INTEGER    :: ij, l
    REAL(rstd) :: wflux_ij, w_ij

    CALL trace_start("compute_caldyn_vert_nh")

#define ETA_DOT(ij) eta_dot(ij,1)
#define WCOV(ij) wcov(ij,1)

    DO l=ll_begin,ll_end
       ! compute the local arrays
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          wflux_ij = .5*(wflux(ij,l)+wflux(ij,l+1))
          w_ij = .5*(W(ij,l)+W(ij,l+1))/mass(ij,l)
          W_etadot(ij,l) = wflux_ij*w_ij
          ETA_DOT(ij) = wflux_ij / mass(ij,l)
          WCOV(ij) = w_ij*(geopot(ij,l+1)-geopot(ij,l))
       ENDDO
       ! add NH term to du
      !DIR$ SIMD
      DO ij=ij_begin,ij_end
          du(ij+u_right,l) = du(ij+u_right,l) &
               - .5*(WCOV(ij+t_right)+WCOV(ij)) &
               *ne_right*(ETA_DOT(ij+t_right)-ETA_DOT(ij))
          du(ij+u_lup,l) = du(ij+u_lup,l) &
               - .5*(WCOV(ij+t_lup)+WCOV(ij)) &
               *ne_lup*(ETA_DOT(ij+t_lup)-ETA_DOT(ij))
          du(ij+u_ldown,l) = du(ij+u_ldown,l) &
               - .5*(WCOV(ij+t_ldown)+WCOV(ij)) &
               *ne_ldown*(ETA_DOT(ij+t_ldown)-ETA_DOT(ij))
       END DO
    ENDDO
    ! add NH terms to dW, dPhi
    ! FIXME : TODO top and bottom
    DO l=ll_beginp1,ll_end ! inner interfaces only
       !DIR$ SIMD
       DO ij=ij_begin,ij_end
          dPhi(ij,l) = dPhi(ij,l) - wflux(ij,l) &
               * (geopot(ij,l+1)-geopot(ij,l-1))/(mass(ij,l-1)+mass(ij,l))
       END DO
    END DO
    DO l=ll_begin,ll_end
       !DIR$ SIMD
       DO ij=ij_begin,ij_end
          dW(ij,l+1) = dW(ij,l+1) + W_etadot(ij,l) ! update inner+top interfaces
          dW(ij,l)   = dW(ij,l)   - W_etadot(ij,l) ! update bottom+inner interfaces
       END DO
    END DO

#undef ETA_DOT
#undef WCOV

    CALL trace_end("compute_caldyn_vert_nh")

  END SUBROUTINE compute_caldyn_vert_NH_manual

END MODULE compute_caldyn_vert_NH_mod