!***************************************************************************************
! lmdz_tendtotke.F90
!*******************
!
! Subroutine that adds a tendency on the TKE created by the
! fluxes of momentum retrieved from the wind speed tendencies
! of the physics. Currently we only account for the contribution
! of the orographic gravity wave drag
!
! The basic concept is the following:
! the TKE equation writes  de/dt = -u'w' du/dz -v'w' dv/dz +g/theta dtheta/dz +......
!
!
! We expect contributions to the term u'w' and v'w' that do not come from turbulent diffusion
! scheme, for instance: orographic gravity wave drag.... These contributions
! need to be accounted for.
! we explicitely calculate the fluxes, integrating the wind speed
!                        tendency from the top of the atmosphere
!
! References: E. Vignon PhD thesis, chapter 7
!             Cheruy et al. 2020, appendix, doi: 10.1029/2019MS002005
!
! contacts: Frederic Hourdin, Etienne Vignon
!
! History:
!---------
! - 1st redaction, Etienne, 15/10/2016
! - add 4 subsurfaces for TKE. Note that the time step disappears in the calculations.
!
!
!**************************************************************************************

!$gpum horizontal klon
MODULE lmdz_gwd_tendtotke
   PRIVATE

   PUBLIC tend_to_tke

CONTAINS

   SUBROUTINE tend_to_tke(dt, klon, klev, nbsrf, is_ave, plev, exner, temp, windu, windv, dt_a, du_a, dv_a, pctsrf, tke)

      USE lmdz_gwd_ini, ONLY: RG, RCPD

      IMPLICIT NONE

! Declarations
!==============

! Inputs
!-------
      REAL, INTENT(IN) ::  dt                   ! Time step [s]
      INTEGER, INTENT(IN)  :: klon, klev ! horizontal and vertical dimensions
      INTEGER, INTENT(IN)  :: nbsrf ! number of subsurfaces
      INTEGER, INTENT(IN)  :: is_ave ! index for mesh-averaged  variables
      REAL, DIMENSION(klon, klev + 1), INTENT(IN) :: plev  ! inter-layer pressure [Pa]
      REAL, DIMENSION(klon, klev), INTENT(IN) :: temp      ! temperature [K], grid-cell average or for a one subsurface
      REAL, DIMENSION(klon, klev), INTENT(IN) :: windu     ! zonal wind [m/s], grid-cell average or for a one subsurface
      REAL, DIMENSION(klon, klev), INTENT(IN) :: windv     ! meridonal wind [m/s], grid-cell average or for a one subsurface
      REAL, DIMENSION(klon, klev), INTENT(IN) :: exner     ! Fonction d'Exner = T/theta
      REAL, DIMENSION(klon, klev), INTENT(IN) :: dt_a      ! Temperature tendency [K], grid-cell average or for a one subsurface
      REAL, DIMENSION(klon, klev), INTENT(IN) :: du_a      ! Zonal wind speed tendency [m/s], grid-cell average or for a one subsurface
      REAL, DIMENSION(klon, klev), INTENT(IN) :: dv_a      ! Meridional wind speed tendency [m/s], grid-cell average or for a one subsurface
      REAL, DIMENSION(klon, nbsrf), INTENT(IN) :: pctsrf   ! fraction of each subsurface [0-1]

! Inputs/Outputs
!---------------
      REAL, DIMENSION(klon, klev + 1, nbsrf + 1), INTENT(INOUT) :: tke       ! Turbulent Kinetic energy [m2/s2], grid-cell average or for a subsurface

! Local
!-------

      INTEGER i, k, isrf                 ! indices
      REAL masse(klon, klev)          ! mass in the layers [kg/m2]
      REAL unsmasse(klon, klev + 1)     ! linear mass in the layers [kg/m2]
      REAL flux_rhotw(klon, klev + 1)   ! flux massique de tempe. pot. rho*u'*theta'
      REAL flux_rhouw(klon, klev + 1)   ! flux massique de quantit?? de mouvement rho*u'*w' [kg/m/s2]
      REAL flux_rhovw(klon, klev + 1)   ! flux massique de quantit?? de mouvement rho*v'*w' [kg/m/s2]
      REAL tendt(klon, klev)        ! new temperature tke tendency [m2/s2/s]
      REAL tendu(klon, klev)        ! new zonal tke tendency [m2/s2/s]
      REAL tendv(klon, klev)        ! new meridonal tke tendency [m2/s2/s]

! First calculations:
!=====================

      unsmasse(:, :) = 0.
      DO k = 1, klev
         masse(:, k) = (plev(:, k) - plev(:, k + 1))/RG
         unsmasse(:, k) = unsmasse(:, k) + 0.5/masse(:, k)
         unsmasse(:, k + 1) = unsmasse(:, k + 1) + 0.5/masse(:, k)
      END DO

      tendu(:, :) = 0.0
      tendv(:, :) = 0.0

! Method 1: Calculation of fluxes using a downward integration
!============================================================

! Flux calculation

      flux_rhotw(:, klev + 1) = 0.
      flux_rhouw(:, klev + 1) = 0.
      flux_rhovw(:, klev + 1) = 0.

      DO k = klev, 1, -1
         flux_rhotw(:, k) = flux_rhotw(:, k + 1) + masse(:, k)*dt_a(:, k)/exner(:, k)
         flux_rhouw(:, k) = flux_rhouw(:, k + 1) + masse(:, k)*du_a(:, k)
         flux_rhovw(:, k) = flux_rhovw(:, k + 1) + masse(:, k)*dv_a(:, k)
      END DO

! TKE update:

      DO k = 2, klev
         tendt(:, k) = -flux_rhotw(:, k)*(exner(:, k) - exner(:, k - 1))*unsmasse(:, k)*RCPD
         tendu(:, k) = -flux_rhouw(:, k)*(windu(:, k) - windu(:, k - 1))*unsmasse(:, k)
         tendv(:, k) = -flux_rhovw(:, k)*(windv(:, k) - windv(:, k - 1))*unsmasse(:, k)
      END DO
      tendt(:, 1) = -flux_rhotw(:, 1)*(exner(:, 1) - 1.)*unsmasse(:, 1)*RCPD
      tendu(:, 1) = -1.*flux_rhouw(:, 1)*windu(:, 1)*unsmasse(:, 1)
      tendv(:, 1) = -1.*flux_rhovw(:, 1)*windv(:, 1)*unsmasse(:, 1)

      DO isrf = 1, nbsrf
         DO k = 1, klev
            DO i = 1, klon
               IF (pctsrf(i, isrf) > 0.) THEN
                  tke(i, k, isrf) = tke(i, k, isrf) + tendu(i, k) + tendv(i, k) + tendt(i, k)
                  tke(i, k, isrf) = max(tke(i, k, isrf), 1.e-10)
               END IF
            END DO
         END DO
      END DO

! Recompute the mesh-averaged TKE
      tke(:, :, is_ave) = 0.
      DO isrf = 1, nbsrf
         DO k = 1, klev
            DO i = 1, klon
               tke(i, k, is_ave) = tke(i, k, is_ave) + pctsrf(i, isrf)*tke(i, k, isrf)
            END DO
         END DO
      END DO

   END SUBROUTINE tend_to_tke

END MODULE lmdz_gwd_tendtotke
