!$gpum horizontal klon
MODULE lmdz_call_gwd

!=====================================================================================================
! lmdz_call_gwd is the interface between the LMDZ physics monitor
! physiq_mod and the routines that compute the drag due to gravity
! waves and subgrid-scale orography (SSO)
! It consists in a sequence of several parameterizations:
! Part I: the effect of orographic gravity waves due to SSO (in fact gravity waves + blocking effect)
! Part II: the effect of non-orographic gravity waves, due to fronts/jets and precipitation ("flott")
! Part III: optional diagnostics of surface torque and angular momentum
! Part IV: effects of mountain drag onto turbulent kinetic energy (TKE)
! contact: F. Lott flott@lmd.ens.fr
!=====================================================================================================

   IMPLICIT NONE

CONTAINS

!===================================================================

   SUBROUTINE call_gwd(klon, klev, nbsrf, is_ter, is_lic, is_ave, abortphy, flag_inhib_tend, itap, JD_cur, JD_ref, JH_cur, &
                       pctsrf, is_sequential, phys_tstep, cell_area, longitude_deg, latitude_deg, pphis, &
                       zstd, zpic, zmea, zval, zsig, zgam, zthe, pplay, paprs, presnivs, rain_fall, snow_fall, &
                       u_beginphys, v_beginphys, rot, tke, wake_delta_tke, &
                       d_u_oro, d_v_oro, d_t_oro, zustr_oro, zvstr_oro, &
                       d_u_lif, d_v_lif, d_t_lif, zustr_lif, zvstr_lif, &
                       d_u_front, d_v_front, zustr_front, zvstr_front, &
                       d_u_precip, d_v_precip, zustr_precip, zvstr_precip, &
                       east_gwstress, west_gwstress, aam, torsfc)

      USE lmdz_gwd_ini, ONLY: ok_strato, ok_orodr, ok_orolf, ok_gwd_rando
      USE lmdz_gwd_ini, ONLY: addtkeoro, smallscales_tkeoro, alphatkeoro
      USE lmdz_gwd_ini, ONLY: zpmm_orodr_t, zstd_orodr_t, zpmm_orolf_t, nm_oro_t
      USE lmdz_gwd_ini, ONLY: ra, rg, rcpd, romega, rkappa

      USE lmdz_gwd_ogwd, ONLY: drag_noro_strato, lift_noro_strato
      USE lmdz_gwd_ogwd_old, ONLY: drag_noro, lift_noro
      USE lmdz_gwd_front, ONLY: acama_gwd_rando
      USE lmdz_gwd_precip, ONLY: flott_gwd_rando

      USE add_phys_tend_mod, ONLY: add_phys_tend, prt_enerbil
      USE phys_local_var_mod, ONLY: u_seri, v_seri, t_seri
      USE lmdz_gwd_tendtotke, ONLY: tend_to_tke

      IMPLICIT NONE

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

! Input variables
!----------------

      INTEGER, INTENT(IN)  :: klon, klev ! number of horizontal and vertical grid points
      INTEGER, INTENT(IN)  :: nbsrf ! number of sub-surfaces
      INTEGER, INTENT(IN)  :: abortphy, flag_inhib_tend, itap ! flags and counter for add_phys_tend
      REAL, INTENT(IN)     :: JD_cur ! current day
      REAL, INTENT(IN)     :: JD_ref ! start day of the run
      REAL, INTENT(IN)     :: JH_cur ! time of the day in seconds
      INTEGER, INTENT(IN)  :: is_ter, is_lic, is_ave ! indices for land and landice subsurfaces and mesh-averaged
      LOGICAL, INTENT(IN)  :: is_sequential ! sequential or parallel model
      REAL, INTENT(IN)  :: phys_tstep ! time step [s]
      REAL, DIMENSION(klon, nbsrf), INTENT(IN)  :: pctsrf ! fraction of each subsurface [0-1]

      REAL, DIMENSION(klon), INTENT(IN)  :: cell_area ! area of the mesh [m2]
      REAL, DIMENSION(klon), INTENT(IN)  :: longitude_deg ! lonfitude of grid points [o]
      REAL, DIMENSION(klon), INTENT(IN)  :: latitude_deg ! latitude of grid points [o]
      REAL, DIMENSION(klon), INTENT(IN)  :: pphis ! surface geopotential [m2/s2]

      REAL, DIMENSION(klon), INTENT(IN)  :: zstd ! std of subgrid-scale orography [m]
      REAL, DIMENSION(klon), INTENT(IN)  :: zpic ! altitude of subgrid-scale orography peaks [m]
      REAL, DIMENSION(klon), INTENT(IN)  :: zmea ! mean altitude of subgrid-scale orography [m]
      REAL, DIMENSION(klon), INTENT(IN)  :: zval ! altitude of subgrid-scale orography valleys [m]
      REAL, DIMENSION(klon), INTENT(IN)  :: zsig ! subgrid-scale orography slope [-]
      REAL, DIMENSION(klon), INTENT(IN)  :: zthe ! subgrid-scale orography small-axis orientation [rad]
      REAL, DIMENSION(klon), INTENT(IN)  :: rain_fall ! rain fall flux at the surface [kg/m2/s]
      REAL, DIMENSION(klon), INTENT(IN)  :: snow_fall ! snowfall flux at the surface [kg/m2/s]

      REAL, DIMENSION(klon, klev), INTENT(IN)  :: pplay ! air pressure [Pa]
      REAL, DIMENSION(klon, klev + 1), INTENT(IN)  :: paprs ! air pressure at bottom layer interface [Pa]
      REAL, DIMENSION(klev), INTENT(IN)  :: presnivs ! equivalent pressure for vertical discretization
      REAL, DIMENSION(klon, klev), INTENT(IN)  :: u_beginphys ! u at the beginning of the physics [m/s]
      REAL, DIMENSION(klon, klev), INTENT(IN)  :: v_beginphys ! v at the beginning of the physics [m/s]
      REAL, DIMENSION(klon, klev), INTENT(IN)  :: rot ! relative vorticity [s-1]

! Inout variables
!-----------------
      REAL, DIMENSION(klon), INTENT(INOUT)  :: zgam ! subgrid-scale orography asymetry parameter [-]
      REAL, DIMENSION(klon, klev + 1, nbsrf + 1), INTENT(INOUT) :: tke ! turbulent kinetic energy [m2/s2]
      REAL, DIMENSION(klon, klev + 1, nbsrf + 1), INTENT(INOUT) :: wake_delta_tke ! turbulent kinetic energy difference between wakes and environment [m2/s2]
      ! tendencies from random non orographic graviy wave processes are inout variables (to account for some process memory).
      REAL, DIMENSION(klon, klev), INTENT(INOUT)  :: d_u_front ! u increment due to gwd generated by fronts [m/s]
      REAL, DIMENSION(klon, klev), INTENT(INOUT)  :: d_u_precip ! u increment due to gwd generated by precipitation [m/s]
      REAL, DIMENSION(klon, klev), INTENT(INOUT)  :: d_v_front ! v increment due to gwd generated by fronts [m/s]
      REAL, DIMENSION(klon, klev), INTENT(INOUT)  :: d_v_precip ! v increment due to gwd generated by precipitation [m/s]
      REAL, DIMENSION(klon, klev), INTENT(INOUT)  :: east_gwstress   ! eastward gravity wave stress [Pa]
      REAL, DIMENSION(klon, klev), INTENT(INOUT)  :: west_gwstress   ! westward gravity wave stress [Pa]

! Output variables
!----------------

      REAL, DIMENSION(klon, klev), INTENT(OUT)  :: d_u_oro ! u increment due to sso drag [m/S]
      REAL, DIMENSION(klon, klev), INTENT(OUT)  :: d_v_oro ! v increment due to sso drag [m/s]
      REAL, DIMENSION(klon, klev), INTENT(OUT)  :: d_t_oro ! t increment due to sso drag [K]

      REAL, DIMENSION(klon, klev), INTENT(OUT)  :: d_u_lif ! u increment due to sso lift [m/s]
      REAL, DIMENSION(klon, klev), INTENT(OUT)  :: d_v_lif ! v increment due to sso lift [m/s]
      REAL, DIMENSION(klon, klev), INTENT(OUT)  :: d_t_lif ! t increment due to sso lift [K]

      REAL, DIMENSION(klon), INTENT(OUT)  :: zustr_oro   ! vertically integrated zonal stress due to sso drag [kg.m/s/m2]
      REAL, DIMENSION(klon), INTENT(OUT)  :: zvstr_oro   ! vertically integrated meridional stress due to sso drag [kg.m/s/m2]

      REAL, DIMENSION(klon), INTENT(OUT)  :: zustr_lif   ! vertically integrated zonal stress due to sso lift [kg.m/s/m2]
      REAL, DIMENSION(klon), INTENT(OUT)  :: zvstr_lif   ! vertically integrated meridional stress due to sso lift [kg.m/s/m2]

      REAL, DIMENSION(klon), INTENT(OUT)  :: zustr_front   ! vertically integrated zonal stress due to fronts gwd param [kg.m/s/m2]
      REAL, DIMENSION(klon), INTENT(OUT)  :: zvstr_front   ! vertically integrated meridional stress due to fronts gwd param [kg.m/s/m2]

      REAL, DIMENSION(klon), INTENT(OUT)  :: zustr_precip   ! vertically integrated zonal stress due to precip. gwd param [kg.m/s/m2]
      REAL, DIMENSION(klon), INTENT(OUT)  :: zvstr_precip   ! vertically integrated meridional stress due to precip. gwd param [kg.m/s/m2]

      REAL, INTENT(OUT) :: aam ! axial wind angular momentum
      REAL, INTENT(OUT) :: torsfc ! total surface torque

! Local variables
!-----------------

      INTEGER :: i, k, igwd
      INTEGER, DIMENSION(klon) :: itest  ! whether the scheme is active (1) or not (0)
      INTEGER, DIMENSION(klon) :: idx  ! indices where the scheme is active
      REAL, DIMENSION(klon) :: nm_oro ! proxy for the number of subgrid-scale mountains
      REAL, DIMENSION(klon) :: zulow, zvlow, zustr_phys, zvstr_phys
      REAL, DIMENSION(klon, klev) :: dt0, dq0, dql0, dqi0, dqbs0, dxt0, dxtl0, dxti0
      REAL, DIMENSION(klon, klev) :: dtadd, duadd, dvadd, exner
      REAL, DIMENSION(klon, klev) :: d_t_oro_tke, d_u_oro_tke, d_v_oro_tke
      REAL, DIMENSION(klon) :: fall

!===================================================================
! Initialisations of local and output variables (as they are not
! always filled in)
!===================================================================

      dt0(:, :) = 0.
      dq0(:, :) = 0.
      dql0(:, :) = 0.
      dqi0(:, :) = 0.
      dqbs0(:, :) = 0.
      dxt0(:, :) = 0.
      dxtl0(:, :) = 0.
      dxti0(:, :) = 0.

      d_u_oro(:, :) = 0.
      d_v_oro(:, :) = 0.
      d_t_oro(:, :) = 0.

      d_u_lif(:, :) = 0.
      d_v_lif(:, :) = 0.
      d_t_lif(:, :) = 0.

      ! DO NOT set tendencies associated with front and precip to 0
      ! as they are inout variables because
      ! the parameterizations considers some "process memory"

      zustr_front(:) = 0.
      zvstr_front(:) = 0.

      zustr_precip(:) = 0.
      zvstr_precip(:) = 0.

      zustr_phys(:) = 0.
      zvstr_phys(:) = 0.

      dtadd(:, :) = 0.
      duadd(:, :) = 0.
      dvadd(:, :) = 0.

      aam = 0.
      torsfc = 0.

      ! calculation of nm_oro
      DO i = 1, klon
         ! nm_oro is a proxy for the number of subgrid scale mountains
         ! -> condition on nm_oro can deactivate SSO params on tilted planar terrains
         !    such as ice sheets (work by V. Wiener).
         ! see Wiener et al. (2025), doi: 10.5194/wcd-6-1605-2025
         ! in such a case, the SSO params should activate only where nm_oro>0 i.e. by setting
         ! nm_oro_t=0.
         nm_oro(i) = zsig(i)*sqrt(cell_area(i)*(pctsrf(i, is_ter) + pctsrf(i, is_lic)))/(4.*MAX(zstd(i), 1.e-8)) - 1.
      END DO

!===================================================================
! Part I: effects of orographic gravity waves
!===================================================================

! I. 1 : drag associated with SSO (that slows down the wind)
!        due to gravity wave breaking and relief blocking effect
!----------------------------------------------------------------

      ! activate or not the param with ok_orodr
      IF (ok_orodr) THEN
         !  we select points where the scheme should be activated (not to treat ocean points mostly)
         igwd = 0
         DO i = 1, klon
            itest(i) = 0
            ! zpmm_orodr_t and zstd_orodr_t are activation thresholds set by F. Lott to
            ! earn computation time but they are not physical (-> set to 0 from CMIP7)
            ! nm_oro_t is a threshold to avoir activate the param over sloping but not mountainous terrains
            IF (((zpic(i) - zmea(i)) .GT. zpmm_orodr_t) .AND. (zstd(i) .GT. zstd_orodr_t) .AND. (nm_oro(i) .GT. nm_oro_t)) THEN
               itest(i) = 1
               igwd = igwd + 1
               idx(igwd) = i
            END IF
         END DO

         IF (ok_strato) THEN

            CALL drag_noro_strato(klon, klev, phys_tstep, 0, paprs, pplay, &
                                  zmea, zstd, zsig, zgam, zthe, zpic, zval, &
                                  igwd, idx, itest, &
                                  t_seri, u_seri, v_seri, &
                                  zulow, zvlow, zustr_oro, zvstr_oro, &
                                  d_t_oro, d_u_oro, d_v_oro)

         ELSE
            ! this routine is becoming obsolete. Do not use it.
            CALL drag_noro(klon, klev, phys_tstep, paprs, pplay, &
                           zmea, zstd, zsig, zgam, zthe, zpic, zval, &
                           igwd, idx, itest, &
                           t_seri, u_seri, v_seri, &
                           zulow, zvlow, zustr_oro, zvstr_oro, &
                           d_t_oro, d_u_oro, d_v_oro)
         END IF

         !  Add tendencies
         !-----------------------------------------------------------------------

#ifdef ISO
         STOP 'appel a add_phys_tend non traite dans call_gwd'
#else
         CALL add_phys_tend(d_u_oro, d_v_oro, d_t_oro, dq0, dql0, dqi0, dqbs0, paprs, 'oro', &
                            abortphy, flag_inhib_tend, itap, 0)
#endif
         CALL prt_enerbil('oro', itap)

      END IF

! I. 2 lift effect due to SSO that changes the direction of the flow (torque effect)
!------------------------------------------------------------------------------------

      ! activate or not the param with ok_orodr
      IF (ok_orolf) THEN

         !  we select points where the scheme should be activated (not to treat ocean points mostly)
         igwd = 0
         DO i = 1, klon
            itest(i) = 0
            ! zpmm_orolf_t is an activation thresholds set by F. Lott to
            ! earn computation time but they are not physical (-> set to 0 from CMIP7)
            ! nm_oro_t is a threshold to avoir activate the param over sloping but not mountainous terrains
            IF (((zpic(i) - zmea(i)) .GT. zpmm_orolf_t) .AND. (nm_oro(i) .GT. nm_oro_t)) THEN
               itest(i) = 1
               igwd = igwd + 1
               idx(igwd) = i
            END IF
         END DO

         IF (ok_strato) THEN

            CALL lift_noro_strato(klon, klev, phys_tstep, paprs, pplay, &
                                  latitude_deg, zmea, zstd, zpic, zgam, zthe, zpic, zval, &
                                  igwd, idx, itest, &
                                  t_seri, u_seri, v_seri, &
                                  zulow, zvlow, zustr_lif, zvstr_lif, &
                                  d_t_lif, d_u_lif, d_v_lif)

         ELSE
            ! this routine is becoming obsolete. Do not use it.
            CALL lift_noro(klon, klev, phys_tstep, paprs, pplay, &
                           latitude_deg, zmea, zstd, zpic, &
                           itest, &
                           t_seri, u_seri, v_seri, &
                           zulow, zvlow, zustr_lif, zvstr_lif, &
                           d_t_lif, d_u_lif, d_v_lif)
         END IF

         !  Add tendencies
         !-----------------------------------------------------------------------

#ifdef ISO
         STOP 'appel a add_phys_tend non traite dans call_gwd'
#else
         CALL add_phys_tend(d_u_lif, d_v_lif, d_t_lif, dq0, dql0, dqi0, dqbs0, paprs, &
                            'lif', abortphy, flag_inhib_tend, itap, 0)
#endif

         CALL prt_enerbil('lif', itap)

      END IF

!===================================================================
! Part II: Non-orographic GW drag
!===================================================================

! II.0 Drag by GWs generated by fronts
!--------------------------------------

      ! stochastic approach following De la Camara & Lott 2015
      IF (ok_gwd_rando) THEN

         CALL acama_gwd_rando(klon, klev, phys_tstep, pplay, presnivs, latitude_deg, t_seri, u_seri, &
                              v_seri, rot, zustr_front, zvstr_front, d_u_front, &
                              d_v_front, east_gwstress, west_gwstress)

         !  Add tendencies
         !-----------------------------------------------------------------------

#ifdef ISO
         STOP 'appel a add_phys_tend non traite dans call_gwd'
#else
         CALL add_phys_tend(d_u_front, d_v_front, dt0, dq0, dql0, dqi0, dqbs0, &
                            paprs, 'gwdfront', abortphy, flag_inhib_tend, itap, 0)
#endif

         CALL prt_enerbil('gwdfront', itap)

         !  Local diagnostics of vertically integrated wind tendencies
         !-----------------------------------------------------------------------
         DO k = 1, klev
            zustr_front(:) = zustr_front(:) + d_u_front(:, k)/phys_tstep &
                             *(paprs(:, k) - paprs(:, k + 1))/rg
            zvstr_front(:) = zvstr_front(:) + d_v_front(:, k)/phys_tstep &
                             *(paprs(:, k) - paprs(:, k + 1))/rg
         END DO

      END IF

! Drag by convective (precip) GWs (stochastic approach) from Lott & Guez 2013
!------------------------------------------------------------------------------

      IF (ok_gwd_rando) THEN

         fall = rain_fall + snow_fall
         CALL flott_gwd_rando(klon, klev, phys_tstep, pplay, presnivs, t_seri, u_seri, v_seri, &
                              fall, zustr_precip, zvstr_precip, &
                              d_u_precip, d_v_precip, east_gwstress, west_gwstress)

         !  Add tendencies
         !-----------------------------------------------------------------------

#ifdef ISO
         STOP 'appel a add_phys_tend non traite dans call_gwd'
#else
         CALL add_phys_tend(d_u_precip, d_v_precip, dt0, dq0, dql0, dqi0, dqbs0, &
                            paprs, 'gwdprecip', abortphy, flag_inhib_tend, itap, 0)
#endif

         CALL prt_enerbil('gwdprecip', itap)

         !  Local diagnostics of vertically integrated wind tendencies
         !-----------------------------------------------------------------------
         DO k = 1, klev
            zustr_precip(:) = zustr_precip(:) + d_u_precip(:, k)/phys_tstep &
                              *(paprs(:, k) - paprs(:, k + 1))/rg
            zvstr_precip(:) = zvstr_precip(:) + d_v_precip(:, k)/phys_tstep &
                              *(paprs(:, k) - paprs(:, k + 1))/rg
         END DO

      END IF

!====================================================================================================
! Part III Diagnostics of mountain-induced torque and angular momentum computation (axial component)
!====================================================================================================

      IF (is_sequential .and. ok_orodr) THEN

         !  Local diagnostics of vertically integrated wind tendencies due to all physics param
         !-------------------------------------------------------------------------------------
         DO k = 1, klev
            DO i = 1, klon
               zustr_phys(i) = zustr_phys(i) + (u_seri(i, k) - u_beginphys(i, k))/phys_tstep* &
                               (paprs(i, k) - paprs(i, k + 1))/rg
               zvstr_phys(i) = zvstr_phys(i) + (v_seri(i, k) - v_beginphys(i, k))/phys_tstep* &
                               (paprs(i, k) - paprs(i, k + 1))/rg
            END DO
         END DO

         !  Mountain-induced torque and angular momentum calculation
         !-------------------------------------------------------------------------------------
         !$gpum nocall
         CALL aaam_bud(27, klon, klev, jD_cur - jD_ref, jH_cur, &
                       ra, rg, romega, &
                       latitude_deg, longitude_deg, pphis, &
                       zustr_oro, zustr_lif, zustr_phys, &
                       zvstr_oro, zvstr_lif, zvstr_phys, &
                       paprs, u_beginphys, v_beginphys, &
                       aam, torsfc)

      END IF

!========================================================================================================
! Part IV TKE tendency associated with gravity waves. Only effect of orographic gravity waves so far
!         see E. Vignon PhD thesis, chapter 7 and Cheruy et al. 2020, appendix, doi: 10.1029/2019MS002005
!========================================================================================================

      ! Choices for addtkeoro:
      !      ** 0 no TKE tendency from orography
      !      ** 1 we include a fraction alphatkeoro of the whole tendency duoro
      !      ** 2 we include a fraction alphatkeoro of the gravity wave part of duoro

      IF (addtkeoro .GT. 0 .AND. ok_orodr) THEN

         IF (addtkeoro .EQ. 1) THEN

            duadd(:, :) = alphatkeoro*d_u_oro(:, :)
            dvadd(:, :) = alphatkeoro*d_v_oro(:, :)

         ELSE IF (addtkeoro .EQ. 2) THEN
            ! this option is obsolete since now
            ! to be removed after tests
            IF (smallscales_tkeoro) THEN
               igwd = 0
               DO i = 1, klon
                  itest(i) = 0
                  ! Here we take into account "all" the subgrid relief (compared to the routine drag_noro_strato
                  ! that activates depending on thresholds as small relief scales can lead to TKE production
                  IF ((zstd(i) .GT. 1.0) .AND. (nm_oro(i) .GT. nm_oro_t)) THEN
                     itest(i) = 1
                     igwd = igwd + 1
                     idx(igwd) = i
                  END IF
               END DO

            ELSE
               igwd = 0
               DO i = 1, klon
                  itest(i) = 0
                 IF (((zpic(i) - zmea(i)) .GT. zpmm_orodr_t) .AND. (zstd(i) .GT. zstd_orodr_t) .AND. (nm_oro(i) .GT. nm_oro_t)) THEN
                     itest(i) = 1
                     igwd = igwd + 1
                     idx(igwd) = i
                  END IF
               END DO
            END IF

            CALL drag_noro_strato(klon, klev, phys_tstep, addtkeoro, paprs, pplay, &
                                  zmea, zstd, zsig, zgam, zthe, zpic, zval, &
                                  igwd, idx, itest, &
                                  t_seri, u_seri, v_seri, &
                                  zulow, zvlow, zustr_oro, zvstr_oro, &
                                  d_t_oro_tke, d_u_oro_tke, d_v_oro_tke)

            zulow(:) = 0.
            zvlow(:) = 0.

            duadd(:, :) = alphatkeoro*d_u_oro_tke(:, :)
            dvadd(:, :) = alphatkeoro*d_v_oro_tke(:, :)
         END IF

         ! TKE update from subgrid temperature and wind tendencies
         !----------------------------------------------------------
         DO k = 1, klev
            exner(:, k) = (pplay(:, k)/paprs(:, 1))**rkappa
         END DO

         CALL tend_to_tke(phys_tstep, klon, klev, nbsrf, is_ave, paprs, exner, t_seri, u_seri, v_seri, dtadd, duadd, dvadd, pctsrf, tke)

         ! Prevent pbl_tke_w from becoming negative
         wake_delta_tke(:, :, :) = max(wake_delta_tke(:, :, :), -tke(:, :, :))

      END IF

   END SUBROUTINE call_gwd

END MODULE lmdz_call_gwd
