MODULE lmdz_atke_turbulence_ini

implicit none

! declaration of constants and parameters

real, save, protected  :: rg, rd, rpi, rcpd, rv, viscom, viscoh        
!$OMP THREADPRIVATE(rg, rd, rpi, rcpd, rv, viscom, viscoh)     

integer, save, protected :: iflag_atke        ! flag that controls options in atke_compute_km_kh
integer, save, protected :: iflag_atke_lmix   ! flag that controls the calculation of mixing length in atke
integer, save, protected :: iflag_num_atke    ! flag that controls the numerical treatment of diffusion coeffiient calculation
!$OMP THREADPRIVATE(iflag_atke, iflag_atke_lmix, iflag_num_atke)

logical, save, protected :: atke_ok_vdiff    ! activate vertical diffusion of TKE or not
logical, save, protected :: atke_ok_virtual  ! account for vapor for flottability
!$OMP THREADPRIVATE(atke_ok_vdiff, atke_ok_virtual)

real, save, protected :: kappa = 0.4         ! Von Karman constant
real, save, protected :: cn                  ! Sm value at Ri=0
real, save, protected :: cinf                ! Sm value at Ri=-Inf
real, save, protected :: ri0                 ! Richardson number near zero to guarantee continuity in slope of Sm (stability function) at Ri=0
real, save, protected :: ri1                 ! Richardson number near zero to guarantee continuity in slope of Pr (Prandlt's number) at Ri=0
real, save, protected :: lmin = 0.01         ! minimum mixing length corresponding to the Kolmogov dissipation scale  in planetary atmospheres (Chen et al 2016, JGR Atmos)
real, save, protected :: ctkes               ! coefficient for surface TKE
real, save, protected :: clmixshear          ! coefficient for mixing length depending on local wind shear
!$OMP THREADPRIVATE(kappa, cn, cinf, ri0, ri1, lmin, ctkes, clmixshear)


! Tunable parameters for the ATKE scheme and their range of values
!!-------------------------------------------------------------------------------------------------------------
real, save, protected :: cepsilon            ! controls the value of the dissipation length scale, range [1.2 - 10]
real, save, protected :: cke                 ! controls the value of the diffusion coefficient of TKE, range [1 - 5]
real, save, protected :: l0                  ! asymptotic mixing length far from the ground [m] (Sun et al 2011, JAMC), range [15 - 75]
real, save, protected :: clmix               ! controls the value of the mixing length in stratified conditions, range [0.1 - 2]
real, save, protected :: ric                 ! critical Richardson number controlling the slope of Sm in stable conditions, range [0.19 - 0.25]
real, save, protected :: smmin               ! minimum value of Sm in very stable conditions (defined here as minsqrt(Ez/Ek)) at large Ri, range [0.025 - 0.1]
real, save, protected :: pr_neut             ! neutral value of the Prandtl number (Ri=0), range [0.7 - 1]
real, save, protected :: pr_slope            ! linear slope of Pr with Ri in the very stable regime, range [3 - 5]
real, save, protected :: cinffac             ! ratio between cinf and cn controlling the convective limit of Sm, range [1.2 - 5.0]
real, save, protected :: pr_asym             ! value of Prandlt in the convective limit(Ri=-Inf), range [0.3 - 0.5]
!$OMP THREADPRIVATE(cepsilon, cke, l0, clmix, ric, smmin, pr_neut, pr_slope, cinffac, pr_asym)
!!-------------------------------------------------------------------------------------------------------------

CONTAINS

SUBROUTINE atke_ini(rg_in, rd_in, rpi_in, rcpd_in, rv_in, viscom_in, viscoh_in)
      !!----------------------------------------------------------------------
      !!                  ***  ROUTINE  atke_ini  ***
      !!
      !! ** Purpose :   Initialization of the atke module and choice of some constants
      !!                Default values correspond to the  'best' configuration
      !!                from tuning on GABLS1 in Vignon et al. 2024, JAMES
      !!----------------------------------------------------------------------

        USE ioipsl_getin_p_mod, ONLY : getin_p                                                                                                      

      ! input arguments (universal constants for planet)
      !-------------------------------------------------
      real, intent(in) :: rg_in, rd_in, rpi_in, rcpd_in, rv_in, viscom_in, viscoh_in
      !!----------------------------------------------------------------------
      
      rg=rg_in          ! gravity acceleration
      rd=rd_in          ! dry gas constant (R/M, R=perfect gas constant and M is the molar mass of the fluid)
      rpi=rpi_in        ! Pi number
      rcpd=rcpd_in      ! cp per unit mass of the gas
      rv=rv_in          ! water vapor constant (for simulations in Earth's atmosphere)
      viscom=viscom_in  ! kinematic molecular viscosity for momentum
      viscoh=viscoh_in  ! kinematic molecular viscosity for heat


      ! Read flag values in .def files
      !-------------------------------

      ! flag that controls options in atke_compute_km_kh
      iflag_atke=1
      CALL getin_p('iflag_atke',iflag_atke)

      ! flag that controls the calculation of mixing length in atke
      iflag_atke_lmix=3
      CALL getin_p('iflag_atke_lmix',iflag_atke_lmix)

      if (iflag_atke .eq. 0 .and. iflag_atke_lmix>0) then
            call abort_physic("atke_turbulence_ini", &
            'stationary scheme must use mixing length formulation not depending on tke', 1)
      endif

      ! activate vertical diffusion of TKE or not
      atke_ok_vdiff=.true.
      CALL getin_p('atke_ok_vdiff',atke_ok_vdiff)

      ! account for vapor for flottability
      atke_ok_virtual=.true.
      CALL getin_p('atke_ok_virtual',atke_ok_virtual)


      ! flag that controls the numerical treatment of diffusion coeffiient calculation
      iflag_num_atke=0
      CALL getin_p('iflag_num_atke',iflag_num_atke)

      ! asymptotic mixing length in neutral conditions [m]
      ! Sun et al 2011, JAMC 
      ! between 10 and 40
      l0=42.5279652116005 
      CALL getin_p('atke_l0',l0)

      ! critical Richardson number 
      ric=0.190537327781655
      CALL getin_p('atke_ric',ric)

      ! constant for tke dissipation calculation
      cepsilon=8.89273387537601
      CALL getin_p('atke_cepsilon',cepsilon)

      ! calculation of cn = Sm value at Ri=0
      ! direct dependance on cepsilon to guarantee Fm=1 (first-order like stability function) at Ri=0
      cn=(1./sqrt(cepsilon))**(2/3)

      ! asymptotic value of Sm for Ri=-Inf
      cinffac=2.0
      CALL getin_p('atke_cinffac',cinffac)
      cinf=cinffac*cn
      if (cinf .le. cn) then
            call abort_physic("atke_turbulence_ini", &
            'cinf cannot be lower than cn', 1)
      endif

      ! coefficient for surface TKE 
      ! following Lenderink & Holtslag 2004, ctkes=(cepsilon)**(2/3)
      ! (provided by limit condition in neutral conditions)
      ctkes=(cepsilon)**(2./3.)

      ! slope of Pr=f(Ri) for stable conditions
      pr_slope=4.67885738180385
      CALL getin_p('atke_pr_slope',pr_slope)
      if (pr_slope .le. 1) then
            call abort_physic("atke_turbulence_ini", &
            'pr_slope has to be greater than 1 for consistency of the tke scheme', 1)
      endif

      ! value of turbulent prandtl number in neutral conditions (Ri=0)
      pr_neut=0.837372701768868
      CALL getin_p('atke_pr_neut',pr_neut)

      ! asymptotic turbulent prandt number value for Ri=-Inf
      pr_asym=0.4
      CALL getin_p('atke_pr_asym',pr_asym)
      if (pr_asym .ge. pr_neut) then
            call abort_physic("atke_turbulence_ini", &
            'pr_asym must be be lower than pr_neut', 1)
      endif

      ! coefficient for mixing length depending on local stratification
      clmix=0.648055235325291
      CALL getin_p('atke_clmix',clmix)

      ! coefficient for mixing length depending on local wind shear
      clmixshear=0.5
      CALL getin_p('atke_clmixshear',clmixshear)

      ! minimum anisotropy coefficient (defined here as minsqrt(Ez/Ek)) at large Ri. 
      ! From Zilitinkevich et al. 2013, it equals sqrt(0.03)~0.17  
      smmin=0.0960838631869678
      CALL getin_p('atke_smmin',smmin)

      ! ratio between the eddy diffusivity coeff for tke wrt that for momentum
      ! default value from Lenderink et al. 2004
      cke=2.47069655134662
      CALL getin_p('atke_cke',cke)

      ! calculation of Ri0 such that continuity in slope of Sm at Ri=0
      ri0=2./rpi*(cinf - cn)*ric/cn

      ! calculation of Ri1 to guarantee continuity in slope of Prandlt number at Ri=0
      ri1 = -2./rpi * (pr_asym - pr_neut)

RETURN

END SUBROUTINE atke_ini

END MODULE  lmdz_atke_turbulence_ini
