module lmdz_blowing_snow_sublim_sedim

contains
subroutine blowing_snow_sublim_sedim(ngrid,nlay,dtime,temp,qv,qb,pplay,paprs,dtemp_bs,dqv_bs,dqb_bs,bsfl,precip_bs)

!==============================================================================
! Routine that calculates the evaporation and sedimentation of blowing snow
! inspired by what is done in lscp_mod
! Etienne Vignon, October 2022
!==============================================================================


use lmdz_blowing_snow_ini, ONLY: iflag_sublim_bs, iflag_sedim_bs, coef_sub_bs,RTT,RD,RG,fallv_bs
use lmdz_blowing_snow_ini, ONLY: qbmin, RCPD, RLSTT, RLMLT, RLVTT, RVTMP2, RV, RPI, tbsmelt, taumeltbs0, rhobs, r_bs
USE lmdz_lscp_tools, ONLY: calc_qsat_ecmwf

implicit none


!++++++++++++++++++++++++++++++++++++++++++++++++++++
! Declarations
!++++++++++++++++++++++++++++++++++++++++++++++++++++

!INPUT
!=====

integer, intent(in)                     :: ngrid,nlay
real, intent(in)                        :: dtime
real, intent(in), dimension(ngrid,nlay) :: temp
real, intent(in), dimension(ngrid,nlay) :: qv
real, intent(in), dimension(ngrid,nlay) :: qb
real, intent(in), dimension(ngrid,nlay) :: pplay
real, intent(in), dimension(ngrid,nlay+1) :: paprs



! OUTPUT
!========


real, intent(out), dimension(ngrid,nlay) :: dtemp_bs
real, intent(out), dimension(ngrid,nlay) :: dqv_bs
real, intent(out), dimension(ngrid,nlay) :: dqb_bs
real, intent(out), dimension(ngrid,nlay+1) :: bsfl
real, intent(out), dimension(ngrid)      :: precip_bs


! LOCAL
!======


integer                                  :: k,i,n
real                                     :: cpd, cpw, dqbsub, maxdqbsub, dqbmelt, zmair
real                                     :: dqsedim,precbs, dqvmelt, zmelt, taumeltbs
real                                     :: maxdqvmelt, rhoair, dz, dqbsedim
real                                     :: delta_p, b_p, a_p, c_p, c_sub, qvsub
real                                     :: p0, T0, Dv, Aprime, Bprime, Ka
real, dimension(ngrid)                   :: ztemp,zqv,zqb,zpres,qsi,dqsi,qsl,dqsl,qzero,sedim
real, dimension(ngrid)                   :: zpaprsup, zpaprsdn, ztemp_up, zqb_up, zvelo
real, dimension(ngrid,nlay)              :: temp_seri, qb_seri, qv_seri

!++++++++++++++++++++++++++++++++++++++++++++++++++
! Initialisation
!++++++++++++++++++++++++++++++++++++++++++++++++++

qzero(:)=0.
dtemp_bs(:,:)=0.
dqv_bs(:,:)=0.
dqb_bs(:,:)=0.
zvelo(:)=0.
sedim(:)=0.
precip_bs(:)=0.
bsfl(:,:)=0.


p0=101325.0    ! ref pressure


DO k=1,nlay
   DO i=1,ngrid
      temp_seri(i,k)=temp(i,k)
      qv_seri(i,k)=qv(i,k)
      qb_seri(i,k)=qb(i,k)
   ENDDO
ENDDO


! Sedimentation scheme
!----------------------

IF (iflag_sedim_bs > 0) THEN
! begin of top-down loop
DO k = nlay, 1, -1
    
    DO i=1,ngrid
        ztemp(i)=temp_seri(i,k)
        zqv(i)=qv_seri(i,k)
        zqb(i)=qb_seri(i,k)
        zpres(i)=pplay(i,k)
        zpaprsup(i)=paprs(i,k+1)
        zpaprsdn(i)=paprs(i,k)
    ENDDO

    ! thermalization of blowing snow precip coming from above     
    IF (k<=nlay-1) THEN

        DO i = 1, ngrid
            zmair=(zpaprsdn(i)-zpaprsup(i))/RG
            ! RVTMP2=rcpv/rcpd-1
            cpd=RCPD*(1.0+RVTMP2*zqv(i))
            cpw=RCPD*RVTMP2
            ! zqb_up: blowing snow mass that has to be thermalized with 
            ! layer's air so that precipitation at the ground has the
            ! same temperature as the lowermost layer
            zqb_up(i) = sedim(i)*dtime/zmair
            ztemp_up(i)=temp_seri(i,k+1)

            ! t(i,k+1)+d_t(i,k+1): new temperature of the overlying layer
            ztemp(i) = ( ztemp_up(i)*zqb_up(i)*cpw + cpd*ztemp(i) ) &
                  / (cpd + zqb_up(i)*cpw)
        ENDDO

     ENDIF

     DO i = 1, ngrid

       rhoair  = zpres(i) / ztemp(i) / RD
       dz      = (zpaprsdn(i)-zpaprsup(i)) / (rhoair*RG)
       ! BS fall velocity assumed to be constant for now
       zvelo(i) = fallv_bs
       ! dqb/dt_sedim=1/rho(sedim/dz - rho w qb/dz)
       ! implicit resolution 
       dqbsedim = (sedim(i)/rhoair/dz*dtime+zqb(i))/(1.+zvelo(i)*dtime/dz) - zqb(i)
       ! flux and dqb update
       zqb(i)=zqb(i)+dqbsedim
       !sedim(i) = sedim(i)-dqbsedim/dtime*rhoair*dz
       sedim(i)=rhoair*zvelo(i)*zqb(i)


       ! variables update:
       bsfl(i,k)=sedim(i)

       qb_seri(i,k) = zqb(i)
       qv_seri(i,k) = zqv(i) 
       temp_seri(i,k) = ztemp(i) 

     ENDDO ! Loop on ngrid


ENDDO ! vertical loop


!surface bs flux
DO i = 1, ngrid
  precip_bs(i) = sedim(i)
ENDDO

ENDIF




!+++++++++++++++++++++++++++++++++++++++++++++++
! Sublimation and melting
!++++++++++++++++++++++++++++++++++++++++++++++
IF (iflag_sublim_bs > 0) THEN

DO k = 1, nlay

    DO i=1,ngrid
        ztemp(i)=temp_seri(i,k)
        zqv(i)=qv_seri(i,k)
        zqb(i)=qb_seri(i,k)
        zpres(i)=pplay(i,k)
        zpaprsup(i)=paprs(i,k+1)
        zpaprsdn(i)=paprs(i,k)
    ENDDO

   ! calulation saturation specific humidity
    CALL CALC_QSAT_ECMWF(ngrid,ztemp(:),qzero(:),zpres(:),RTT,2,.false.,qsi(:),dqsi(:))


    DO i = 1, ngrid

          rhoair  = zpres(i) / ztemp(i) / RD
          dz      = (zpaprsdn(i)-zpaprsup(i)) / (rhoair*RG)
          ! BS fall velocity assumed to be constant for now
          zvelo(i) = fallv_bs


          IF (ztemp(i) > RTT) THEN

             ! if temperature is positive, we assume that part of the blowing snow 
             ! already present  melts and evaporates with a typical time 
             ! constant taumeltbs

             taumeltbs=taumeltbs0*exp(-max(0.,(ztemp(i)-RTT)/(tbsmelt-RTT)))
             dqvmelt=min(zqb(i),-1.*zqb(i)*(exp(-dtime/taumeltbs)-1.))
             maxdqvmelt= max(RCPD*(1.0+RVTMP2*(zqv(i)+zqb(i)))*(ztemp(i)-RTT)/(RLMLT+RLVTT),0.)
             dqvmelt=min(dqvmelt,maxdqvmelt)
             ! qv update, melting + evaporation
             zqv(i) = zqv(i) + dqvmelt
             ! temp update melting + evaporation
             ztemp(i) = ztemp(i) - dqvmelt * (RLMLT+RLVTT)/RCPD/(1.0+RVTMP2*(zqv(i)+zqb(i)))
             ! qb update melting + evaporation
             zqb(i)=zqb(i)-dqvmelt

          ELSE
              ! negative celcius temperature     
              ! Sublimation scheme   


              ! Sublimation formulation for ice crystals from Pruppacher & Klett, Rutledge & Hobbs 1983
              ! assuming monodispered crystal distrib
              ! dqb/dt_sub=-coef_sub_bs*(1-qv/qsi)*nc*8*r_bs/(Aprime+Bprime)
              ! assuming Mi=rhobs*4/3*pi*r_bs**3 
              ! rhoair qb=nc*Mi -> nc=rhoair qb/Mi
              ! dqb/dt_sub=-coef_sub_bs*(1-qv/qsi)*6*rhoair*qb/(rhobs*pi*r_bs**2)/(Aprime+Bprime)
              ! dqb/dt_sub=-c_sub(1-qv/qsi)*qb
              ! c_sub=coef_sub_bs*6*rhoair/(rhobs*pi*r_bs**2)/(Aprime+Bprime)

              ! Note the strong coupling between specific contents of water vapor and blowing snow during sublimation
              ! equations dqv/dt_sub and dqb/dt_sub must be solved jointly to prevent irrealistic increase of water vapor
              ! at typical physics time steps
              ! we thus solve the differential equation using an implicit scheme for both qb and qv
 
              ! we do not consider deposition, only sublimation
              IF (zqv(i) < qsi(i)) THEN
                 rhoair=zpres(i)/ztemp(i)/RD
                 Dv=0.0001*0.211*(p0/zpres(i))*((ztemp(i)/RTT)**1.94) ! water vapor diffusivity in air, SI
                 Ka=(5.69+0.017*(ztemp(i)-RTT))*1.e-5*100.*4.184                ! thermal conductivity of the air, SI
                 Aprime=RLSTT/Ka/ztemp(i)*(RLSTT/RV/ztemp(i) -1.)
                 Bprime=1./(rhoair*Dv*qsi(i)) 
                 c_sub=coef_sub_bs*6.*rhoair/(rhobs*RPI*r_bs**2)/(Aprime+Bprime)
                 c_p=-zqb(i)
                 b_p=1.+c_sub*dtime-c_sub*dtime/qsi(i)*zqb(i)-c_sub*dtime/qsi(i)*zqv(i)
                 a_p=c_sub*dtime/qsi(i)
                 delta_p=(b_p**2)-4.*a_p*c_p  
                 dqbsub=(-b_p+sqrt(delta_p))/(2.*a_p) - zqb(i)       
                 dqbsub = MIN(0.0,MAX(dqbsub,-zqb(i)))
                 ! Sublimation limit: we ensure that the whole mesh does not exceed saturation wrt ice
                 maxdqbsub = MAX(0.0, qsi(i)-zqv(i))
                 dqbsub = MAX(dqbsub,-maxdqbsub)
              ELSE
                 dqbsub=0.
              ENDIF

              ! vapor, temperature, precip fluxes update following sublimation
              zqv(i) = zqv(i) - dqbsub
              zqb(i) = zqb(i) + dqbsub
              ztemp(i) = ztemp(i) + dqbsub*RLSTT/RCPD/(1.0+RVTMP2*(zqv(i)+zqb(i)))

          ENDIF

          ! if qb<qbmin, sublimate or melt and evaporate qb
          ! see Gerber et al. 2023, JGR Atmos for the choice of qbmin

          IF (zqb(i) < qbmin) THEN
              zqv(i) = zqv(i)+zqb(i)
              IF (ztemp(i) < RTT) THEN
                 ztemp(i) = ztemp(i) - zqb(i) * RLSTT/RCPD/(1.0+RVTMP2*(zqv(i)))
              ELSE
                 ztemp(i) = ztemp(i) - zqb(i) * (RLVTT+RLMLT)/RCPD/(1.0+RVTMP2*(zqv(i)))
              ENDIF
              zqb(i)=0.
          ENDIF

     ! variables update
     temp_seri(i,k)=ztemp(i)
     qv_seri(i,k)=zqv(i)
     qb_seri(i,k)=zqb(i)
     ENDDO
ENDDO

ENDIF


! OUTPUTS
!++++++++++

! 3D variables
DO k=1,nlay
   DO i=1,ngrid
        dqb_bs(i,k) = qb_seri(i,k) - qb(i,k)
        dqv_bs(i,k) = qv_seri(i,k) - qv(i,k)
        dtemp_bs(i,k) = temp_seri(i,k) - temp(i,k)
   ENDDO
ENDDO



return

end subroutine blowing_snow_sublim_sedim
end module lmdz_blowing_snow_sublim_sedim
