MODULE lmdz_ratqs_multi

!=============================================
! A FAIRE :
! Traiter le probleme de USE lmdz_lscp_tools, ONLY: CALC_QSAT_ECMWF
!=============================================

!=============================================
! module containing subroutines that take 
! into account the effect of convection, orography,
! surface heterogeneities and subgrid-scale
! turbulence on ratqs, i.e. on the width of the
! total water subgrid distribution.
!=============================================

USE yoethf_mod_h
    IMPLICIT NONE

! Include
!=============================================


      CONTAINS


!========================================================================    
SUBROUTINE ratqs_inter(klon,klev,iflag_ratqs,pdtphys,paprs, &
           ratqsbas, wake_deltaq, wake_s, q_seri,qtc_cv, sigt_cv,     &
           fm_therm,entr_therm,detr_therm,detrain_cv,fm_cv,fqd,fqcomp,sigd, &
           ratqs_inter_,sigma_qtherm)

USE lmdz_ratqs_ini, ONLY : a_ratqs_cv,tau_var,fac_tau,tau_cumul,a_ratqs_wake, dqimpl
USE lmdz_ratqs_ini, ONLY : RG
USE lmdz_ratqs_ini, ONLY : povariance, var_conv
USE lmdz_thermcell_dq,  ONLY : thermcell_dq

implicit none

!========================================================================
! L. d'Alen??on, 25/02/2021
! Cette subroutine calcule une valeur de ratqsbas interactive  
! Elle est appel??e par la subroutine ratqs lorsque iflag_ratqs = 11. 
!========================================================================

! Declarations
! Input
integer,intent(in) :: klon,klev,iflag_ratqs
real,intent(in) :: pdtphys,ratqsbas
real, dimension(klon,klev+1),intent(in) :: paprs
real, dimension(klon,klev),intent(in) :: wake_deltaq, q_seri,qtc_cv, sigt_cv
real, dimension(klon),intent(in) :: wake_s
real, dimension(klon,klev+1),intent(in) :: fm_therm
real, dimension(klon,klev),intent(in) :: entr_therm,detr_therm,detrain_cv,fm_cv,fqd,fqcomp
real, dimension(klon),intent(in) :: sigd

! Output
real, dimension(klon,klev),intent(inout) :: ratqs_inter_

! local
LOGICAL :: klein = .false.
LOGICAL :: klein_conv = .true.
REAL :: taup0 = 70000
REAL :: taudp = 500
integer :: lev_out=10
REAL, DIMENSION (klon,klev) :: zmasse,entr0,detr0,detraincv,dqp,detrain_p,q0,qd0,tau_diss
REAL, DIMENSION (klon,klev+1) :: fm0
integer i,k
real, dimension(klon,klev) :: wake_dq

real, dimension(klon) :: max_sigd, max_dqconv,max_sigt
real, dimension(klon,klev) :: zoa,zocarrea,pdocarreadj,pocarre,po,pdoadj,varq_therm
real, dimension(klon,klev) :: var_moy, var_var, var_desc_th,var_det_conv,var_desc_prec,var_desc_conv,sigma_qtherm

lev_out=0.

print*,'ratqs_inter'

!-----------------------------------------------------------------------
!   Calcul des masses
!-----------------------------------------------------------------------

      do k=1,klev
         zmasse(:,k)=(paprs(:,k)-paprs(:,k+1))/RG
      enddo
!-------------------------------------------------------------------------
!  Caclul du terme de d??trainement de la variance pour les thermiques
!-------------------------------------------------------------------------

! initialisations 


      do k=1,klev
         do i=1,klon
            tau_diss(i,k)=tau_var +0.5*fac_tau*tau_var*(tanh((taup0-paprs(i,k))/taudp) + 1.)
         enddo
      enddo
       
                    
      
      entr0(:,:) = entr_therm(:,:) 
      fm0(:,:) = fm_therm(:,:)  
      detr0(:,:) = detr_therm(:,:) 

! calcul du carr?? de l'humidit?? sp??cifique et circulation dans les thermiques
      po(:,:) = q_seri(:,:)
      call thermcell_dq(klon,klev,dqimpl,pdtphys,fm0,entr0,zmasse,  &
     &                   po,pdoadj,zoa,lev_out)
      do k=1,klev
         do i=1,klon
            pocarre(i,k)=po(i,k)*po(i,k) + povariance(i,k)
         enddo
      enddo
      call thermcell_dq(klon,klev,dqimpl,pdtphys,fm0,entr0,zmasse,  &
      &                   pocarre,pdocarreadj,zocarrea,lev_out) 


! variance de l'humidit?? sp??cifique totale dans les thermiques     
      do k=1,klev
         do i=1,klon      
            varq_therm(i,k)=zocarrea(i,k)-zoa(i,k)*zoa(i,k)
         enddo
      enddo

! calcul des termes sources de la variance avec thermiques et convection profonde (voir Klein 2005 par exemple)
      do k=1,klev
         do i=1,klon      
            var_moy(i,k) = detr0(i,k)*((zoa(i,k)-po(i,k))**2)/zmasse(i,k)
            var_var(i,k) = detr0(i,k)*(varq_therm(i,k)-povariance(i,k))/zmasse(i,k)
            var_det_conv(i,k) =  a_ratqs_cv*(detrain_cv(i,k)/zmasse(i,k))
            if (sigd(i).ne.0) then
               var_desc_prec(i,k) = sigd(i)*(1-sigd(i))*(fqd(i,k)*tau_cumul/sigd(i))**2/tau_cumul
            else
               var_desc_prec(i,k) = 0
            endif
         enddo
      enddo

      do k=1,klev-1
         do i=1,klon      
            var_desc_th(i,k) = fm0(i,k+1)*povariance(i,k+1)/zmasse(i,k) -  &
               fm0(i,k)*povariance(i,k)/zmasse(i,k)
            var_desc_conv(i,k) = ((povariance(i,k+1)-povariance(i,k))*(fm_cv(i,k)/zmasse(i,k)))
         enddo
      enddo
      var_desc_th(:,klev) = var_desc_th(:,klev-1)
      var_desc_conv(:,klev) = var_desc_conv(:,klev-1)
      
      if (klein) then
         do k=1,klev-1
            do i=1,klon
              qd0(:,:) = 0.0
              if (sigd(i).ne.0) then
                qd0(i,k) = fqd(i,k)*tau_cumul/sigd(i) 
              endif
            enddo
         enddo
         do k=1,klev-1
            do i=1,klon      
               povariance(i,k)= (var_moy(i,k) + var_var(i,k) + var_desc_th(i,k) +  &
               var_det_conv(i,k) +  var_desc_prec(i,k)  &   
                + var_desc_conv(i,k))*pdtphys + povariance(i,k)
               povariance(i,k)= povariance(i,k)*exp(-pdtphys/tau_diss(i,k))
            enddo
         enddo
         povariance(:,klev) = povariance(:,klev-1)
         
      else ! calcul direct     
         qd0(:,:) = 0.0
         q0(:,:) = 0.0
         do k=1,klev-1
            do i=1,klon
               if (sigd(i).ne.0) then    ! termes de variance par accumulation
                 qd0(i,k) = fqd(i,k)*tau_cumul/sigd(i) 
               endif
               if (sigt_cv(i,k).ne.0) then
                 q0(i,k) = fqcomp(i,k)*tau_cumul/sigt_cv(i,k)
               endif
            enddo
         enddo
         do k=1,klev-1
            do i=1,klon      
               povariance(i,k)= (pdocarreadj(i,k)-2.*po(i,k)*pdoadj(i,k) +  &
               a_ratqs_cv*(sigt_cv(i,k)*(1-sigt_cv(i,k))*q0(i,k)**2/tau_cumul + var_desc_prec(i,k) +  &
               var_desc_conv(i,k)))*pdtphys + povariance(i,k)
               povariance(i,k)=povariance(i,k)*exp(-pdtphys/tau_diss(i,k))
            enddo
         enddo
         povariance(:,klev) = povariance(:,klev-1)
!         fqd(:,:)=sigt_cv(:,:)*(1-sigt_cv(:,:))*q0(:,:)**2/tau_cumul 
      endif

!-------------------------------------------------------------------------
!  Caclul du ratqs_inter_
!-------------------------------------------------------------------------

      do k=1,klev
        do i=1,klon
           if(q_seri(i,k).ge.1E-7) then
               ratqs_inter_(i,k) = abs(povariance(i,k))**0.5/q_seri(i,k)    
               sigma_qtherm(i,k) = abs(varq_therm(i,k))**0.5     ! sigma dans les thermiques
           else 
               ratqs_inter_(i,k) = 0.  
               sigma_qtherm(i,k) = 0.
           endif
        enddo
      enddo
      
return
END SUBROUTINE ratqs_inter

!------------------------------------------------------------------
SUBROUTINE ratqs_oro(klon,klev,pctsrf,zstd,qsat,temp,pplay,paprs,ratqs_oro_)

! Etienne Vignon, November 2021: effect of subgrid orography on ratqs

USE lmdz_ratqs_ini, ONLY : RG,RV,RD,RLSTT,RLVTT,RTT,nbsrf,is_lic,is_ter

IMPLICIT NONE

! Declarations
!--------------

! INPUTS

INTEGER, INTENT(IN) :: klon                       ! number of horizontal grid points
INTEGER, INTENT(IN) :: klev                       ! number of vertical layers
REAL, DIMENSION(klon,nbsrf) :: pctsrf
REAL, DIMENSION(klon,klev), INTENT(IN) :: qsat    ! saturation specific humidity [kg/kg]
REAL, DIMENSION(klon), INTENT(IN) :: zstd    ! sub grid orography standard deviation
REAL, DIMENSION(klon,klev), INTENT(IN) :: temp    ! air temperature [K]
REAL, DIMENSION(klon,klev), INTENT(IN) :: pplay    ! air pressure, layer's center [Pa]
REAL, DIMENSION(klon,klev+1), INTENT(IN) :: paprs    ! air pressure, lower inteface [Pa]

! OUTPUTS

REAL, DIMENSION(klon,klev), INTENT(out) :: ratqs_oro_ ! ratqs profile due to subgrid orography


! LOCAL

INTEGER :: i,k
REAL, DIMENSION(klon) :: orogradT,xsi0
REAL, DIMENSION (klon,klev) :: zlay
REAL :: Lvs, temp0


! Calculation of the near-surface temperature gradient along the topography
!--------------------------------------------------------------------------

! at the moment, we fix it at a constant value (moist adiab. lapse rate)

orogradT(:)=-6.5/1000. ! K/m

! Calculation of near-surface surface ratqs
!-------------------------------------------

DO i=1,klon
    temp0=temp(i,1)
    IF (temp0 .LT. RTT) THEN
        Lvs=RLSTT
    ELSE
        Lvs=RLVTT
    ENDIF
    xsi0(i)=zstd(i)*ABS(orogradT(i))*Lvs/temp0/temp0/RV
    ratqs_oro_(i,1)=xsi0(i)
END DO

! Vertical profile of ratqs assuming an exponential decrease with height
!------------------------------------------------------------------------
     
! calculation of geop. height AGL        
zlay(:,1)= RD*temp(:,1)/(0.5*(paprs(:,1)+pplay(:,1))) &
           *(paprs(:,1)-pplay(:,1))/RG

DO k=2,klev
   DO i = 1, klon
      zlay(i,k)= zlay(i,k-1)+RD*0.5*(temp(i,k-1)+temp(i,k)) &
               /paprs(i,k)*(pplay(i,k-1)-pplay(i,k))/RG
               
      ratqs_oro_(i,k)=MAX(0.0,pctsrf(i,is_ter)*xsi0(i)*exp(-zlay(i,k)/MAX(zstd(i),1.)))   
    END DO
END DO




END SUBROUTINE ratqs_oro

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

SUBROUTINE ratqs_hetero(klon,klev,pctsrf,s_pblh,t2m,q2m,temp,q,pplay,paprs,ratqs_hetero_)

! Etienne Vignon, November 2021
! Effect of subgrid surface heterogeneities on ratqs

USE lmdz_lscp_tools, ONLY: CALC_QSAT_ECMWF

USE lmdz_ratqs_ini, ONLY : RG,RD,RTT,nbsrf

IMPLICIT NONE

! INPUTS


INTEGER, INTENT(IN) :: klon                       ! number of horizontal grid points
INTEGER, INTENT(IN) :: klev                       ! number of vertical layers
REAL, DIMENSION(klon)                   :: s_pblh ! height of the planetary boundary layer(HPBL)
REAL, DIMENSION(klon,nbsrf)             :: pctsrf ! Fractional cover of subsurfaces
REAL, DIMENSION(klon,nbsrf), INTENT(IN) :: t2m    ! 2m temperature for each tile [K]
REAL, DIMENSION(klon,nbsrf), INTENT(IN) :: q2m    ! 2m specific humidity for each tile [kg/kg]
REAL, DIMENSION(klon,klev), INTENT(IN) :: temp    ! air temperature [K]
REAL, DIMENSION(klon,klev), INTENT(IN) :: q       ! specific humidity [kg/kg]
REAL, DIMENSION(klon,klev), INTENT(IN) :: pplay   ! air pressure, layer's center [Pa]
REAL, DIMENSION(klon,klev+1), INTENT(IN) :: paprs ! air pressure, lower inteface [Pa]

! OUTPUTS

REAL, DIMENSION(klon,klev), INTENT(out) :: ratqs_hetero_ ! ratsq profile due to surface heterogeneities


INTEGER :: i,k,nsrf
REAL, DIMENSION(klon) :: xsi0, ratiom, qsat2m, dqsatdT
REAL, DIMENSION (klon,klev) :: zlay



! Calculation of near-surface surface ratqs
!-------------------------------------------

    
    ratiom(:)=0.
    xsi0(:)=0.
    
    DO nsrf=1,nbsrf
    CALL CALC_QSAT_ECMWF(klon,t2m(:,nsrf),q2m(:,nsrf),paprs(:,1),RTT,0,.false.,qsat2m,dqsatdT)
    ratiom(:)=ratiom(:)+pctsrf(:,nsrf)*(q2m(:,nsrf)/qsat2m(:))
    xsi0(:)=xsi0(:)+pctsrf(:,nsrf)*((q2m(:,nsrf)/qsat2m(:)-ratiom(:))**2)
    END DO
    
    xsi0(:)=sqrt(xsi0(:))/(ratiom(:)+1E-6)



! Vertical profile of ratqs assuming an exponential decrease with height
!------------------------------------------------------------------------
        
! calculation of geop. height AGL

zlay(:,1)= RD*temp(:,1)/(0.5*(paprs(:,1)+pplay(:,1))) &
           *(paprs(:,1)-pplay(:,1))/RG
ratqs_hetero_(:,1)=xsi0(:)

DO k=2,klev
   DO i = 1, klon
      zlay(i,k)= zlay(i,k-1)+RD*0.5*(temp(i,k-1)+temp(i,k)) &
               /paprs(i,k)*(pplay(i,k-1)-pplay(i,k))/RG
               
      ratqs_hetero_(i,k)=MAX(xsi0(i)*exp(-zlay(i,k)/(s_pblh(i)+1.0)),0.0)   
    END DO
END DO

END SUBROUTINE ratqs_hetero

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

SUBROUTINE ratqs_tke(klon,klev,pdtphys,temp,q,qsat,pplay,paprs,omega,tke,tke_dissip,lmix,wprime,ratqs_tke_)

! References:
!
! Etienne Vignon: effect of subgrid turbulence on ratqs
!
! Field, P.R., Hill, A., Furtado, K., Korolev, A., 2014b. Mixed-phase clouds in a turbulent environment. Part
! 2: analytic treatment. Q. J. R. Meteorol. Soc. 21, 2651???2663. https://doi.org/10.1002/qj.2175.
! 
! Furtado, K., Field, P.R., Boutle, I.A., Morcrette, C.R., Wilkinson, J., 2016. A physically-based, subgrid
! parametrization for the production and maintenance of mixed-phase clouds in a general circulation
! model. J. Atmos. Sci. 73, 279???291. https://doi.org/10.1175/JAS-D-15-0021.

USE lmdz_ratqs_ini, ONLY : RG,RV,RD,RCPD,RLSTT,RLVTT,RTT

IMPLICIT NONE

! INPUTS

INTEGER, INTENT(IN) :: klon                             ! number of horizontal grid points
INTEGER, INTENT(IN) :: klev                             ! number of vertical layers
REAL, INTENT(IN) :: pdtphys                             ! physics time step [s]
REAL, DIMENSION(klon,klev), INTENT(IN) :: temp          ! air temperature [K]
REAL, DIMENSION(klon,klev), INTENT(IN) :: q             ! specific humidity [kg/kg]
REAL, DIMENSION(klon,klev), INTENT(IN) :: qsat          ! saturation specific humidity [kg/kg]
REAL, DIMENSION(klon,klev), INTENT(IN) :: pplay         ! air pressure, layer's center [Pa]
REAL, DIMENSION(klon,klev+1), INTENT(IN) :: paprs       ! air pressure, lower inteface [Pa]
REAL, DIMENSION(klon,klev), INTENT(IN) :: omega       ! air pressure, lower inteface [Pa]
REAL, DIMENSION(klon,klev+1), INTENT(IN) :: tke         ! Turbulent Kinetic Energy [m2/s2]
REAL, DIMENSION(klon,klev+1), INTENT(IN) :: tke_dissip  ! Turbulent Kinetic Energy Dissipation rate [m2/s3]
REAL, DIMENSION(klon,klev+1), INTENT(IN) :: lmix  ! Turbulent mixing length
REAL, DIMENSION(klon,klev+1), INTENT(IN) :: wprime      ! Turbulent vertical velocity scale [m/s]

! OUTPUTS

REAL, DIMENSION(klon,klev), INTENT(out) :: ratqs_tke_  ! ratsq profile due to subgrid TKE

! LOCAL
INTEGER :: i, k
REAL :: AA, DD, NW, AAprime, VARLOG,rho,Lvs,taue,lhomo,dissmin,maxvarlog
REAL, DIMENSION(klon,klev) :: sigmaw,w
REAL, PARAMETER :: C0=10.0
REAL, PARAMETER :: lmin=0.001
REAL, PARAMETER :: ratqsmin=1E-6
REAL, PARAMETER :: ratqsmax=0.5


! Calculation of large scale and turbulent vertical velocities
!---------------------------------------------------------------

DO k=1,klev
    DO i=1,klon
        rho=pplay(i,k)/temp(i,k)/RD
        w(i,k)=-rho*RG*omega(i,k)
        sigmaw(i,k)=0.5*(wprime(i,k+1)+wprime(i,k)) ! turbulent vertical velocity at the middle of model layers. 
    END DO
END DO

! Calculation of ratqs
!---------------------------------------------------------------
ratqs_tke_(:,1)=ratqsmin ! set to a very low value to avoid division by 0 in order parts
                        ! of the code
DO k=2,klev ! we start from second model level since TKE is not defined at k=1
    DO i=1,klon

       IF (temp(i,k) .LT. RTT) THEN
           Lvs=RLSTT
       ELSE
           Lvs=RLVTT
       ENDIF
       dissmin=0.01*(0.5*(tke(i,k)+tke(i,k+1))/pdtphys)
       maxvarlog=LOG(1.0+ratqsmax**2)! to prevent ratqs from exceeding an arbitrary threshold value 
       AA=RG*(Lvs/(RCPD*temp(i,k)*temp(i,k)*RV) - 1./(RD*temp(i,k)))
       lhomo=MAX(0.5*(lmix(i,k)+lmix(i,k+1)),lmin)
       taue=(lhomo*lhomo/MAX(0.5*(tke_dissip(i,k)+tke_dissip(i,k+1)),dissmin))**(1./3) ! Fields et al. 2014
       DD=1.0/taue
       NW=(sigmaw(i,k)**2)*SQRT(2./(C0*MAX(0.5*(tke_dissip(i,k)+tke_dissip(i,k+1)),dissmin)))
       AAprime=AA*NW
       VARLOG=AAprime/2./DD
       VARLOG=MIN(VARLOG,maxvarlog)
       ratqs_tke_(i,k)=SQRT(MAX(EXP(VARLOG)-1.0,ratqsmin))
       END DO
END DO
END SUBROUTINE ratqs_tke

END MODULE lmdz_ratqs_multi
