MODULE caldyn_kernels_hevi_mod
  USE icosa
  USE transfert_mod
  USE caldyn_kernels_base_mod
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: compute_theta, compute_pvort_only, compute_caldyn_fast, compute_caldyn_slow

CONTAINS

  SUBROUTINE compute_theta(ps,theta_rhodz, rhodz,theta)
    USE icosa
    USE disvert_mod, ONLY : mass_dak, mass_dbk, caldyn_eta, eta_mass, ptop
    USE exner_mod
    USE trace
    USE omp_para
    IMPLICIT NONE
    REAL(rstd),INTENT(IN)  :: ps(iim*jjm)
    REAL(rstd),INTENT(IN)  :: theta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: theta(iim*jjm,llm)
    INTEGER :: ij,l
    REAL(rstd) :: m
    CALL trace_start("compute_theta")  

    IF(caldyn_eta==eta_mass) THEN ! Compute mass & theta
       DO l = ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             IF(DEC) THEN  ! ps is actually Ms
                m = mass_dak(l)+(ps(ij)*g+ptop)*mass_dbk(l)
             ELSE
                m = mass_dak(l)+ps(ij)*mass_dbk(l)
             END IF
             rhodz(ij,l) = m/g
             theta(ij,l) = theta_rhodz(ij,l)/rhodz(ij,l)
          ENDDO
       ENDDO
    ELSE ! Compute only theta
       DO l = ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             theta(ij,l) = theta_rhodz(ij,l)/rhodz(ij,l)
          ENDDO
       ENDDO
    END IF

    CALL trace_end("compute_theta")
  END SUBROUTINE compute_theta

  SUBROUTINE compute_pvort_only(u,rhodz,qu,qv)
    USE icosa
    USE exner_mod
    USE trace
    USE omp_para
    IMPLICIT NONE
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: qu(iim*3*jjm,llm)
    REAL(rstd),INTENT(OUT) :: qv(iim*2*jjm,llm)

    INTEGER :: ij,l
    REAL(rstd) :: etav,hv,radius_m2

    CALL trace_start("compute_pvort_only")  
!!! Compute shallow-water potential vorticity
    radius_m2=radius**(-2)
    DO l = ll_begin,ll_end
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          IF(DEC) THEN
             etav= 1./Av(ij+z_up)*(  ne_rup  * u(ij+u_rup,l)         &
                                   + ne_left * u(ij+t_rup+u_left,l)  &
                                   - ne_lup  * u(ij+u_lup,l) )                               

             hv =   Riv2(ij,vup)          * rhodz(ij,l)           &
                  + Riv2(ij+t_rup,vldown) * rhodz(ij+t_rup,l)     &
                  + Riv2(ij+t_lup,vrdown) * rhodz(ij+t_lup,l)
             qv(ij+z_up,l) = ( etav+fv(ij+z_up) )/hv

             etav = 1./Av(ij+z_down)*(  ne_ldown * u(ij+u_ldown,l)          &
                                      + ne_right * u(ij+t_ldown+u_right,l)  &
                                      - ne_rdown * u(ij+u_rdown,l) )
             hv =   Riv2(ij,vdown)        * rhodz(ij,l)          &
                  + Riv2(ij+t_ldown,vrup) * rhodz(ij+t_ldown,l)  &
                  + Riv2(ij+t_rdown,vlup) * rhodz(ij+t_rdown,l)
             qv(ij+z_down,l) =( etav+fv(ij+z_down) )/hv
          ELSE
             etav= 1./Av(ij+z_up)*(  ne_rup  * u(ij+u_rup,l)        * de(ij+u_rup)         &
                                   + ne_left * u(ij+t_rup+u_left,l) * de(ij+t_rup+u_left)  &
                                   - ne_lup  * u(ij+u_lup,l)        * de(ij+u_lup) )                               

             hv =   Riv2(ij,vup)          * rhodz(ij,l)           &
                  + Riv2(ij+t_rup,vldown) * rhodz(ij+t_rup,l)     &
                  + Riv2(ij+t_lup,vrdown) * rhodz(ij+t_lup,l)
             qv(ij+z_up,l) = ( etav+fv(ij+z_up) )/hv

             etav = 1./Av(ij+z_down)*(  ne_ldown * u(ij+u_ldown,l)          * de(ij+u_ldown)          &
                                      + ne_right * u(ij+t_ldown+u_right,l)  * de(ij+t_ldown+u_right)  &
                                      - ne_rdown * u(ij+u_rdown,l)          * de(ij+u_rdown) )
             hv =   Riv2(ij,vdown)        * rhodz(ij,l)          &
                  + Riv2(ij+t_ldown,vrup) * rhodz(ij+t_ldown,l)  &
                  + Riv2(ij+t_rdown,vlup) * rhodz(ij+t_rdown,l)
             qv(ij+z_down,l) =( etav+fv(ij+z_down) )/hv
          END IF
       ENDDO

       !DIR$ SIMD
       DO ij=ij_begin,ij_end
          qu(ij+u_right,l) = 0.5*(qv(ij+z_rdown,l)+qv(ij+z_rup,l))  
          qu(ij+u_lup,l) = 0.5*(qv(ij+z_up,l)+qv(ij+z_lup,l))  
          qu(ij+u_ldown,l) = 0.5*(qv(ij+z_ldown,l)+qv(ij+z_down,l))  
       END DO

    ENDDO

    CALL trace_end("compute_pvort_only")

  END SUBROUTINE compute_pvort_only

  SUBROUTINE compute_caldyn_fast(tau,u,rhodz,theta,pk,geopot, du)
    USE icosa
    USE disvert_mod
    USE exner_mod
    USE trace
    USE omp_para
    IMPLICIT NONE
    REAL(rstd), INTENT(IN) :: tau ! "solve" u-tau*du/dt = rhs
    REAL(rstd),INTENT(INOUT)  :: u(iim*3*jjm,llm)    ! prognostic "velocity"
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm)  ! potential temperature
    REAL(rstd),INTENT(INOUT) :: pk(iim*jjm,llm) ! Exner function
    REAL(rstd),INTENT(IN)  :: geopot(iim*jjm,llm+1)    ! geopotential
    REAL(rstd),INTENT(OUT) :: du(iim*3*jjm,llm)
    REAL(rstd) :: berni(iim*jjm,llm)  ! Bernoulli function

    INTEGER :: i,j,ij,l
    REAL(rstd) :: due_right, due_lup, due_ldown

    CALL trace_start("compute_caldyn_fast")
    
    ! Compute bernouilli term
    IF(boussinesq) THEN
       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             berni(ij,l) = pk(ij,l)
             ! from now on pk contains the vertically-averaged geopotential
             pk(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1))
          ENDDO
       ENDDO

    ELSE ! compressible

       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1))
          ENDDO
       ENDDO

    END IF ! Boussinesq/compressible

!!! u:=u+tau*du, du = gradients of Bernoulli and Exner functions
    DO l=ll_begin,ll_end
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          due_right =  0.5*(theta(ij,l)+theta(ij+t_right,l))                  &
                          *(ne_right*pk(ij,l)   +ne_left*pk(ij+t_right,l))    &
                         +  ne_right*berni(ij,l)+ne_left*berni(ij+t_right,l)
          due_lup = 0.5*(theta(ij,l)+theta(ij+t_lup,l))                       &
                       *(ne_lup*pk(ij,l)   +ne_rdown*pk(ij+t_lup,l))          &
                      +  ne_lup*berni(ij,l)+ne_rdown*berni(ij+t_lup,l)
          due_ldown = 0.5*(theta(ij,l)+theta(ij+t_ldown,l))                   &
                         *(ne_ldown*pk(ij,l)   +ne_rup*pk(ij+t_ldown,l))      &
                        +  ne_ldown*berni(ij,l)+ne_rup*berni(ij+t_ldown,l)
          IF(.NOT.DEC) THEN
             due_right = due_right/de(ij+u_right)
             due_lup   = due_lup/de(ij+u_lup)
             due_ldown = due_ldown/de(ij+u_ldown)
          END IF
          du(ij+u_right,l) = due_right
          du(ij+u_lup,l)   = due_lup
          du(ij+u_ldown,l) = due_ldown
          u(ij+u_right,l) = u(ij+u_right,l) + tau*due_right
          u(ij+u_lup,l)   = u(ij+u_lup,l)   + tau*due_lup
          u(ij+u_ldown,l) = u(ij+u_ldown,l) + tau*due_ldown
       ENDDO
    ENDDO

    CALL trace_end("compute_caldyn_fast")

  END SUBROUTINE compute_caldyn_fast

  SUBROUTINE compute_caldyn_slow(u,rhodz,qu,theta, hflux,convm, dtheta_rhodz, du)
    USE icosa
    USE disvert_mod
    USE exner_mod
    USE trace
    USE omp_para
    IMPLICIT NONE
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)    ! prognostic "velocity"
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)  :: qu(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm)  ! potential temperature

    REAL(rstd),INTENT(OUT) :: hflux(iim*3*jjm,llm) ! hflux in kg/s
    REAL(rstd),INTENT(OUT) :: convm(iim*jjm,llm)  ! mass flux convergence
    REAL(rstd),INTENT(OUT) :: dtheta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: du(iim*3*jjm,llm)

    REAL(rstd) :: cor_NT(iim*jjm,llm)  ! NT coriolis force u.(du/dPhi)
    REAL(rstd) :: urel(3*iim*jjm,llm)  ! relative velocity
    REAL(rstd) :: Ftheta(3*iim*jjm,llm) ! theta flux
    REAL(rstd) :: berni(iim*jjm,llm)  ! Bernoulli function

    INTEGER :: ij,l
    REAL(rstd) :: uu_right, uu_lup, uu_ldown

    CALL trace_start("compute_caldyn_slow")

    DO l = ll_begin, ll_end
!!!  Compute mass and theta fluxes
       IF (caldyn_conserv==energy) CALL test_message(req_qu) 
       !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)
          IF(DEC) THEN
             uu_right= uu_right*le(ij+u_right)/de(ij+u_right)
             uu_lup  = uu_lup  *le(ij+u_lup)/de(ij+u_lup)
             uu_ldown= uu_ldown*le(ij+u_ldown)/de(ij+u_ldown)
          ELSE
             uu_right= uu_right*le(ij+u_right)
             uu_lup  = uu_lup  *le(ij+u_lup)
             uu_ldown= uu_ldown*le(ij+u_ldown)
          END IF
          hflux(ij+u_right,l)=uu_right
          hflux(ij+u_lup,l)  =uu_lup
          hflux(ij+u_ldown,l)=uu_ldown
          Ftheta(ij+u_right,l)=0.5*(theta(ij,l)+theta(ij+t_right,l))*uu_right
          Ftheta(ij+u_lup,l)=0.5*(theta(ij,l)+theta(ij+t_lup,l))*uu_lup
          Ftheta(ij+u_ldown,l)=0.5*(theta(ij,l)+theta(ij+t_ldown,l))*uu_ldown
       ENDDO

!!! compute horizontal divergence of fluxes
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          ! convm = -div(mass flux), sign convention as in Ringler et al. 2012, eq. 21
          convm(ij,l)= -1./Ai(ij)*(ne_right*hflux(ij+u_right,l)   +  &
               ne_rup*hflux(ij+u_rup,l)       +  &  
               ne_lup*hflux(ij+u_lup,l)       +  &  
               ne_left*hflux(ij+u_left,l)     +  &  
               ne_ldown*hflux(ij+u_ldown,l)   +  &  
               ne_rdown*hflux(ij+u_rdown,l))    

          ! signe ? attention d (rho theta dz)
          ! dtheta_rhodz =  -div(flux.theta)
          dtheta_rhodz(ij,l)=-1./Ai(ij)*(ne_right*Ftheta(ij+u_right,l)    +  &
               ne_rup*Ftheta(ij+u_rup,l)        +  &  
               ne_lup*Ftheta(ij+u_lup,l)        +  &  
               ne_left*Ftheta(ij+u_left,l)      +  &  
               ne_ldown*Ftheta(ij+u_ldown,l)    +  &  
               ne_rdown*Ftheta(ij+u_rdown,l))
       ENDDO

    END DO

!!! Compute potential vorticity (Coriolis) contribution to du

    SELECT CASE(caldyn_conserv)
    CASE(energy) ! energy-conserving TRiSK

       CALL wait_message(req_qu)

       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             uu_right = &
                  wee(ij+u_right,1,1)*hflux(ij+u_rup,l)*(qu(ij+u_right,l)+qu(ij+u_rup,l))+                             &
                  wee(ij+u_right,2,1)*hflux(ij+u_lup,l)*(qu(ij+u_right,l)+qu(ij+u_lup,l))+                             &
                  wee(ij+u_right,3,1)*hflux(ij+u_left,l)*(qu(ij+u_right,l)+qu(ij+u_left,l))+                           &
                  wee(ij+u_right,4,1)*hflux(ij+u_ldown,l)*(qu(ij+u_right,l)+qu(ij+u_ldown,l))+                         &
                  wee(ij+u_right,5,1)*hflux(ij+u_rdown,l)*(qu(ij+u_right,l)+qu(ij+u_rdown,l))+                         & 
                  wee(ij+u_right,1,2)*hflux(ij+t_right+u_ldown,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_ldown,l))+         &
                  wee(ij+u_right,2,2)*hflux(ij+t_right+u_rdown,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_rdown,l))+         &
                  wee(ij+u_right,3,2)*hflux(ij+t_right+u_right,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_right,l))+         &
                  wee(ij+u_right,4,2)*hflux(ij+t_right+u_rup,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_rup,l))+             &
                  wee(ij+u_right,5,2)*hflux(ij+t_right+u_lup,l)*(qu(ij+u_right,l)+qu(ij+t_right+u_lup,l))
             uu_lup = &
                  wee(ij+u_lup,1,1)*hflux(ij+u_left,l)*(qu(ij+u_lup,l)+qu(ij+u_left,l)) +                        &
                  wee(ij+u_lup,2,1)*hflux(ij+u_ldown,l)*(qu(ij+u_lup,l)+qu(ij+u_ldown,l)) +                      &
                  wee(ij+u_lup,3,1)*hflux(ij+u_rdown,l)*(qu(ij+u_lup,l)+qu(ij+u_rdown,l)) +                      &
                  wee(ij+u_lup,4,1)*hflux(ij+u_right,l)*(qu(ij+u_lup,l)+qu(ij+u_right,l)) +                      &
                  wee(ij+u_lup,5,1)*hflux(ij+u_rup,l)*(qu(ij+u_lup,l)+qu(ij+u_rup,l)) +                          & 
                  wee(ij+u_lup,1,2)*hflux(ij+t_lup+u_right,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_right,l)) +          &
                  wee(ij+u_lup,2,2)*hflux(ij+t_lup+u_rup,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_rup,l)) +              &
                  wee(ij+u_lup,3,2)*hflux(ij+t_lup+u_lup,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_lup,l)) +              &
                  wee(ij+u_lup,4,2)*hflux(ij+t_lup+u_left,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_left,l)) +            &
                  wee(ij+u_lup,5,2)*hflux(ij+t_lup+u_ldown,l)*(qu(ij+u_lup,l)+qu(ij+t_lup+u_ldown,l))
             uu_ldown = &
                  wee(ij+u_ldown,1,1)*hflux(ij+u_rdown,l)*(qu(ij+u_ldown,l)+qu(ij+u_rdown,l)) +                      &
                  wee(ij+u_ldown,2,1)*hflux(ij+u_right,l)*(qu(ij+u_ldown,l)+qu(ij+u_right,l)) +                      &
                  wee(ij+u_ldown,3,1)*hflux(ij+u_rup,l)*(qu(ij+u_ldown,l)+qu(ij+u_rup,l)) +                          &
                  wee(ij+u_ldown,4,1)*hflux(ij+u_lup,l)*(qu(ij+u_ldown,l)+qu(ij+u_lup,l)) +                          &
                  wee(ij+u_ldown,5,1)*hflux(ij+u_left,l)*(qu(ij+u_ldown,l)+qu(ij+u_left,l)) +                        & 
                  wee(ij+u_ldown,1,2)*hflux(ij+t_ldown+u_lup,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_lup,l)) +          &
                  wee(ij+u_ldown,2,2)*hflux(ij+t_ldown+u_left,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_left,l)) +        &
                  wee(ij+u_ldown,3,2)*hflux(ij+t_ldown+u_ldown,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_ldown,l)) +      &
                  wee(ij+u_ldown,4,2)*hflux(ij+t_ldown+u_rdown,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_rdown,l)) +      &
                  wee(ij+u_ldown,5,2)*hflux(ij+t_ldown+u_right,l)*(qu(ij+u_ldown,l)+qu(ij+t_ldown+u_right,l))
             IF(DEC) THEN
                du(ij+u_right,l) = .5*uu_right
                du(ij+u_lup,l)   = .5*uu_lup
                du(ij+u_ldown,l)   = .5*uu_ldown
             ELSE
                du(ij+u_right,l) = .5*uu_right/de(ij+u_right)
                du(ij+u_lup,l)   = .5*uu_lup  /de(ij+u_lup)
                du(ij+u_ldown,l) = .5*uu_ldown/de(ij+u_ldown)
             END IF
          ENDDO
       ENDDO

    CASE(enstrophy) ! enstrophy-conserving TRiSK

       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             uu_right = &
                  wee(ij+u_right,1,1)*hflux(ij+u_rup,l)+                           &
                  wee(ij+u_right,2,1)*hflux(ij+u_lup,l)+                           &
                  wee(ij+u_right,3,1)*hflux(ij+u_left,l)+                          &
                  wee(ij+u_right,4,1)*hflux(ij+u_ldown,l)+                         &
                  wee(ij+u_right,5,1)*hflux(ij+u_rdown,l)+                         & 
                  wee(ij+u_right,1,2)*hflux(ij+t_right+u_ldown,l)+                 &
                  wee(ij+u_right,2,2)*hflux(ij+t_right+u_rdown,l)+                 &
                  wee(ij+u_right,3,2)*hflux(ij+t_right+u_right,l)+                 &
                  wee(ij+u_right,4,2)*hflux(ij+t_right+u_rup,l)+                   &
                  wee(ij+u_right,5,2)*hflux(ij+t_right+u_lup,l)
             uu_lup = &
                  wee(ij+u_lup,1,1)*hflux(ij+u_left,l)+                        &
                  wee(ij+u_lup,2,1)*hflux(ij+u_ldown,l)+                       &
                  wee(ij+u_lup,3,1)*hflux(ij+u_rdown,l)+                       &
                  wee(ij+u_lup,4,1)*hflux(ij+u_right,l)+                       &
                  wee(ij+u_lup,5,1)*hflux(ij+u_rup,l)+                         & 
                  wee(ij+u_lup,1,2)*hflux(ij+t_lup+u_right,l)+                 &
                  wee(ij+u_lup,2,2)*hflux(ij+t_lup+u_rup,l)+                   &
                  wee(ij+u_lup,3,2)*hflux(ij+t_lup+u_lup,l)+                   &
                  wee(ij+u_lup,4,2)*hflux(ij+t_lup+u_left,l)+                  &
                  wee(ij+u_lup,5,2)*hflux(ij+t_lup+u_ldown,l)
             uu_ldown = &
                  wee(ij+u_ldown,1,1)*hflux(ij+u_rdown,l)+                      &
                  wee(ij+u_ldown,2,1)*hflux(ij+u_right,l)+                      &
                  wee(ij+u_ldown,3,1)*hflux(ij+u_rup,l)+                        &
                  wee(ij+u_ldown,4,1)*hflux(ij+u_lup,l)+                        &
                  wee(ij+u_ldown,5,1)*hflux(ij+u_left,l)+                       & 
                  wee(ij+u_ldown,1,2)*hflux(ij+t_ldown+u_lup,l)+                &
                  wee(ij+u_ldown,2,2)*hflux(ij+t_ldown+u_left,l)+               &
                  wee(ij+u_ldown,3,2)*hflux(ij+t_ldown+u_ldown,l)+              &
                  wee(ij+u_ldown,4,2)*hflux(ij+t_ldown+u_rdown,l)+              &
                  wee(ij+u_ldown,5,2)*hflux(ij+t_ldown+u_right,l)
             IF(DEC) THEN
                du(ij+u_right,l) = qu(ij+u_right,l)*uu_right
                du(ij+u_lup,l)   = qu(ij+u_lup,l)  *uu_lup
                du(ij+u_ldown,l) = qu(ij+u_ldown,l)*uu_ldown
             ELSE
                du(ij+u_right,l) = qu(ij+u_right,l)*uu_right/de(ij+u_right)
                du(ij+u_lup,l)   = qu(ij+u_lup,l)  *uu_lup  /de(ij+u_lup)
                du(ij+u_ldown,l) = qu(ij+u_ldown,l)*uu_ldown/de(ij+u_ldown)
             END IF
          ENDDO
       ENDDO

    CASE DEFAULT
       STOP
    END SELECT

    ! Compute bernouilli term = Kinetic Energy
    le_de(:) = le(:)/de(:) ! FIXME - make sure le_de is what we expect
    DO l=ll_begin,ll_end
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          IF(DEC) THEN
             berni(ij,l) = &
                  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 )  
          ELSE
             berni(ij,l) = &
                  1/(4*Ai(ij))*(le(ij+u_right)*de(ij+u_right)*u(ij+u_right,l)**2 +    &
                  le(ij+u_rup)*de(ij+u_rup)*u(ij+u_rup,l)**2 +          &
                  le(ij+u_lup)*de(ij+u_lup)*u(ij+u_lup,l)**2 +          &
                  le(ij+u_left)*de(ij+u_left)*u(ij+u_left,l)**2 +       &
                  le(ij+u_ldown)*de(ij+u_ldown)*u(ij+u_ldown,l)**2 +    &
                  le(ij+u_rdown)*de(ij+u_rdown)*u(ij+u_rdown,l)**2 )  
          END IF
       ENDDO
    ENDDO

!!! Add gradients of Bernoulli and Exner functions to du
    DO l=ll_begin,ll_end
       !DIR$ SIMD
       DO ij=ij_begin,ij_end
          IF(DEC) THEN
             du(ij+u_right,l) = du(ij+u_right,l)                       &
                  + ne_right*berni(ij,l)+ne_left*berni(ij+t_right,l)
             du(ij+u_lup,l) = du(ij+u_lup,l)                           &
                  + ne_lup*berni(ij,l)+ne_rdown*berni(ij+t_lup,l)
             du(ij+u_ldown,l) = du(ij+u_ldown,l)                       &
                  + ne_ldown*berni(ij,l)+ne_rup*berni(ij+t_ldown,l)
          ELSE
             du(ij+u_right,l) = du(ij+u_right,l) + 1/de(ij+u_right)            &
                  * ( ne_right*berni(ij,l)+ne_left*berni(ij+t_right,l) )
             du(ij+u_lup,l) = du(ij+u_lup,l) +  1/de(ij+u_lup)                 &
                  * ( ne_lup*berni(ij,l)+ne_rdown*berni(ij+t_lup,l) )
             du(ij+u_ldown,l) = du(ij+u_ldown,l) + 1/de(ij+u_ldown)            &
                  * ( ne_ldown*berni(ij,l)+ne_rup*berni(ij+t_ldown,l) )
          END IF
       END DO
    ENDDO

    CALL trace_end("compute_caldyn_slow")

  END SUBROUTINE compute_caldyn_slow

END MODULE caldyn_kernels_hevi_mod