MODULE compute_caldyn_slow_hydro_mod
  USE grid_param, ONLY : llm
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: compute_caldyn_slow_hydro

CONTAINS

  SUBROUTINE compute_caldyn_slow_hydro(u,rhodz,hv, hflux,Kv,du, zero)
    USE icosa
    USE omp_para, ONLY : ll_begin, ll_end
    USE caldyn_vars_mod
    LOGICAL, INTENT(IN) :: zero
    REAL(rstd),INTENT(IN)  :: u(3*iim*jjm,llm)    ! prognostic "velocity"
    REAL(rstd),INTENT(IN)  :: Kv(2*iim*jjm,llm)   ! kinetic energy at vertices
    REAL(rstd),INTENT(IN)  :: hv(2*iim*jjm,llm)   ! height/mass averaged to vertices
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: hflux(3*iim*jjm,llm) ! hflux in kg/s
    REAL(rstd),INTENT(INOUT) :: du(3*iim*jjm,llm)
    
    REAL(rstd) :: berni(iim*jjm,llm)  ! Bernoulli function
    REAL(rstd) :: berni1(iim*jjm)  ! Bernoulli function
    REAL(rstd) :: uu_right, uu_lup, uu_ldown, ke, uu
    INTEGER :: ij,l

    CALL trace_start("compute_caldyn_slow_hydro")

    IF(dysl_slow_hydro) THEN

#define BERNI(ij,l) berni(ij,l)
#include "../kernels_hex/caldyn_slow_hydro.k90"
#undef BERNI

     ELSE

#define BERNI(ij) berni1(ij)

    DO l = ll_begin, ll_end
       !  Compute mass fluxes
       IF (caldyn_conserv==conserv_energy) CALL test_message(req_qu) 

       IF(caldyn_kinetic==kinetic_trisk) THEN
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             uu_right=0.5*(rhodz(ij,l)+rhodz(ij+t_right,l))*u(ij+u_right,l)
             uu_lup=0.5*(rhodz(ij,l)+rhodz(ij+t_lup,l))*u(ij+u_lup,l)
             uu_ldown=0.5*(rhodz(ij,l)+rhodz(ij+t_ldown,l))*u(ij+u_ldown,l)
             uu_right= uu_right*le_de(ij+u_right)
             uu_lup  = uu_lup  *le_de(ij+u_lup)
             uu_ldown= uu_ldown*le_de(ij+u_ldown)
             hflux(ij+u_right,l)=uu_right
             hflux(ij+u_lup,l)  =uu_lup
             hflux(ij+u_ldown,l)=uu_ldown
          ENDDO
       ELSE ! mass flux deriving from consistent kinetic energy
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             uu_right=0.5*(hv(ij+z_rup,l)+hv(ij+z_rdown,l))*u(ij+u_right,l)
             uu_lup=0.5*(hv(ij+z_up,l)+hv(ij+z_lup,l))*u(ij+u_lup,l)
             uu_ldown=0.5*(hv(ij+z_ldown,l)+hv(ij+z_down,l))*u(ij+u_ldown,l)
             uu_right= uu_right*le_de(ij+u_right)
             uu_lup  = uu_lup  *le_de(ij+u_lup)
             uu_ldown= uu_ldown*le_de(ij+u_ldown)
             hflux(ij+u_right,l)=uu_right
             hflux(ij+u_lup,l)  =uu_lup
             hflux(ij+u_ldown,l)=uu_ldown
          ENDDO
       END IF

       ! Compute Bernoulli=kinetic energy
       IF(caldyn_kinetic==kinetic_trisk) THEN
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             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 )  
          ENDDO
       ELSE
          !DIR$ SIMD
          DO ij=ij_begin,ij_end
             BERNI(ij) = Riv(ij,vup)   *Kv(ij+z_up,l)    + &
                         Riv(ij,vlup)  *Kv(ij+z_lup,l)   + &
                         Riv(ij,vldown)*Kv(ij+z_ldown,l) + &
                         Riv(ij,vdown) *Kv(ij+z_down,l)  + &
                         Riv(ij,vrdown)*Kv(ij+z_rdown,l) + &
                         Riv(ij,vrup)  *Kv(ij+z_rup,l)
          END DO
       END IF
       ! Compute du=-grad(Bernoulli)
       IF(zero) THEN
          !DIR$ SIMD
          DO ij=ij_begin,ij_end
             du(ij+u_right,l) = ne_right*(BERNI(ij)-BERNI(ij+t_right))
             du(ij+u_lup,l)   = ne_lup*(BERNI(ij)-BERNI(ij+t_lup))
             du(ij+u_ldown,l) = ne_ldown*(BERNI(ij)-BERNI(ij+t_ldown))
          END DO
       ELSE
          !DIR$ SIMD
          DO ij=ij_begin,ij_end
             du(ij+u_right,l) = du(ij+u_right,l) + &
                  ne_right*(BERNI(ij)-BERNI(ij+t_right))
             du(ij+u_lup,l)   = du(ij+u_lup,l) + &
                  ne_lup*(BERNI(ij)-BERNI(ij+t_lup))
             du(ij+u_ldown,l) = du(ij+u_ldown,l) + &
                  ne_ldown*(BERNI(ij)-BERNI(ij+t_ldown))
          END DO
       END IF
    END DO

#undef BERNI

    END IF ! dysl
    CALL trace_end("compute_caldyn_slow_hydro")    
  END SUBROUTINE compute_caldyn_slow_hydro

END MODULE compute_caldyn_slow_hydro_mod