MODULE climb_qbs_mod

  ! Module to solve the verctical diffusion of blowing snow;

  USE dimphy

  IMPLICIT NONE
  SAVE
  PRIVATE
  PUBLIC :: climb_qbs_down, climb_qbs_up

  REAL, DIMENSION(:, :), ALLOCATABLE :: gamaqbs
  !$OMP THREADPRIVATE(gamaqbs)
  REAL, DIMENSION(:, :), ALLOCATABLE :: Ccoef_QBS, Dcoef_QBS
  !$OMP THREADPRIVATE(Ccoef_QBS, Dcoef_QBS)
  REAL, DIMENSION(:), ALLOCATABLE :: Acoef_QBS, Bcoef_QBS
  !$OMP THREADPRIVATE(Acoef_QBS, Bcoef_QBS)
  REAL, DIMENSION(:, :), ALLOCATABLE :: Kcoefqbs
  !$OMP THREADPRIVATE(Kcoefqbs)

CONTAINS

  !****************************************************************************************

  SUBROUTINE climb_qbs_down(knon, coefqbs, paprs, pplay, &
          delp, temp, qbs, dtime, &
          Ccoef_QBS_out, Dcoef_QBS_out, &
          Kcoef_qbs_out, gama_qbs_out, &
          Acoef_QBS_out, Bcoef_QBS_out)

    ! This routine calculates recursivly the coefficients C and D
    ! for the quantity X=[QBS] in equation X(k) = C(k) + D(k)*X(k-1), where k is
    ! the index of the vertical layer.

    ! Input arguments
    !****************************************************************************************
    USE lmdz_compbl, ONLY: iflag_pbl, iflag_pbl_split, iflag_order2_sollw, ifl_pbltree
    USE lmdz_yomcst

    IMPLICIT NONE
    INTEGER, INTENT(IN) :: knon
    REAL, DIMENSION(klon, klev), INTENT(IN) :: coefqbs
    REAL, DIMENSION(klon, klev), INTENT(IN) :: pplay
    REAL, DIMENSION(klon, klev + 1), INTENT(IN) :: paprs
    REAL, DIMENSION(klon, klev), INTENT(IN) :: delp
    REAL, DIMENSION(klon, klev), INTENT(IN) :: temp
    REAL, DIMENSION(klon, klev), INTENT(IN) :: qbs
    REAL, INTENT(IN) :: dtime

    ! Output arguments
    !****************************************************************************************
    REAL, DIMENSION(klon), INTENT(OUT) :: Acoef_QBS_out
    REAL, DIMENSION(klon), INTENT(OUT) :: Bcoef_QBS_out

    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Ccoef_QBS_out
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Dcoef_QBS_out
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Kcoef_qbs_out
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: gama_qbs_out

    ! Local variables
    !****************************************************************************************
    LOGICAL, SAVE :: first = .TRUE.
    !$OMP THREADPRIVATE(first)
    REAL, DIMENSION(klon) :: psref
    REAL :: delz, pkh
    INTEGER :: k, i, ierr
    !****************************************************************************************
    ! 1)
    ! Allocation at first time step only

    !****************************************************************************************

    IF (first) THEN
      first = .FALSE.
      ALLOCATE(Ccoef_QBS(klon, klev), STAT = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in allloc Ccoef_QBS, ierr=', ierr

      ALLOCATE(Dcoef_QBS(klon, klev), STAT = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in allloc Dcoef_QBS, ierr=', ierr

      ALLOCATE(Acoef_QBS(klon), Bcoef_QBS(klon), STAT = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in allloc Acoef_BS and Bcoef_BS, ierr=', ierr

      ALLOCATE(Kcoefqbs(klon, klev), STAT = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in allloc Kcoefqbs, ierr=', ierr

      ALLOCATE(gamaqbs(1:klon, 2:klev), STAT = ierr)
      IF (ierr /= 0) PRINT*, ' pb in allloc gamaqbs, ierr=', ierr

    END IF

    !****************************************************************************************
    ! 2)
    ! Definition of the coeficient K

    !****************************************************************************************
    Kcoefqbs(:, :) = 0.0
    DO k = 2, klev
      DO i = 1, knon
        Kcoefqbs(i, k) = &
                coefqbs(i, k) * RG * RG * dtime / (pplay(i, k - 1) - pplay(i, k)) &
                        * (paprs(i, k) * 2 / (temp(i, k) + temp(i, k - 1)) / RD)**2
      ENDDO
    ENDDO

    !****************************************************************************************
    ! 3)
    ! Calculation of gama for "Q" and "H"

    !****************************************************************************************
    !   surface pressure is used as reference
    psref(:) = paprs(:, 1)

    !   definition of gama
    IF (iflag_pbl == 1) THEN
      gamaqbs(:, :) = 0.0

      ! conversion de gama
      DO k = 2, klev
        DO i = 1, knon
          delz = RD * (temp(i, k - 1) + temp(i, k)) / &
                  2.0 / RG / paprs(i, k) * (pplay(i, k - 1) - pplay(i, k))
          pkh = (psref(i) / paprs(i, k))**RKAPPA

          ! convertie gradient verticale de contenu en neige soufflee en difference de neige soufflee entre centre de couches
          gamaqbs(i, k) = gamaqbs(i, k) * delz
        ENDDO
      ENDDO

    ELSE
      gamaqbs(:, :) = 0.0
    ENDIF


    !****************************************************************************************
    ! 4)
    ! Calculte the coefficients C and D for specific content of blowing snow, qbs

    !****************************************************************************************

    CALL calc_coef_qbs(knon, Kcoefqbs(:, :), gamaqbs(:, :), delp(:, :), qbs(:, :), &
            Ccoef_QBS(:, :), Dcoef_QBS(:, :), Acoef_QBS, Bcoef_QBS)


    !****************************************************************************************
    ! 5)
    ! Return the first layer in output variables

    !****************************************************************************************
    Acoef_QBS_out = Acoef_QBS
    Bcoef_QBS_out = Bcoef_QBS

    !****************************************************************************************
    ! 6)
    ! If Pbl is split, return also the other layers in output variables

    !****************************************************************************************
    IF (mod(iflag_pbl_split, 10) >=1) THEN
      DO k = 1, klev
        DO i = 1, klon
          Ccoef_QBS_out(i, k) = Ccoef_QBS(i, k)
          Dcoef_QBS_out(i, k) = Dcoef_QBS(i, k)
          Kcoef_qbs_out(i, k) = Kcoefqbs(i, k)
          IF (k==1) THEN
            gama_qbs_out(i, k) = 0.
          ELSE
            gama_qbs_out(i, k) = gamaqbs(i, k)
          ENDIF
        ENDDO
      ENDDO
      !!!
    ENDIF  ! (mod(iflag_pbl_split,2) .ge.1)
    !!!

  END SUBROUTINE climb_qbs_down

  !****************************************************************************************

  SUBROUTINE calc_coef_qbs(knon, Kcoef, gama, delp, X, Ccoef, Dcoef, Acoef, Bcoef)

    ! Calculate the coefficients C and D in : X(k) = C(k) + D(k)*X(k-1)
    ! where X is QQBS, and k the vertical level k=1,klev
    USE lmdz_yomcst
    IMPLICIT NONE
    ! Input arguments
    !****************************************************************************************
    INTEGER, INTENT(IN) :: knon
    REAL, DIMENSION(klon, klev), INTENT(IN) :: Kcoef, delp
    REAL, DIMENSION(klon, klev), INTENT(IN) :: X
    REAL, DIMENSION(klon, 2:klev), INTENT(IN) :: gama

    ! Output arguments
    !****************************************************************************************
    REAL, DIMENSION(klon), INTENT(OUT) :: Acoef, Bcoef
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Ccoef, Dcoef

    ! Local variables
    !****************************************************************************************
    INTEGER :: k, i
    REAL :: buf

    !****************************************************************************************
    ! Niveau au sommet, k=klev

    !****************************************************************************************
    Ccoef(:, :) = 0.0
    Dcoef(:, :) = 0.0

    DO i = 1, knon
      buf = delp(i, klev) + Kcoef(i, klev)

      Ccoef(i, klev) = (X(i, klev) * delp(i, klev) - Kcoef(i, klev) * gama(i, klev)) / buf
      Dcoef(i, klev) = Kcoef(i, klev) / buf
    END DO


    !****************************************************************************************
    ! Niveau  (klev-1) <= k <= 2

    !****************************************************************************************

    DO k = (klev - 1), 2, -1
      DO i = 1, knon
        buf = delp(i, k) + Kcoef(i, k) + Kcoef(i, k + 1) * (1. - Dcoef(i, k + 1))
        Ccoef(i, k) = (X(i, k) * delp(i, k) + Kcoef(i, k + 1) * Ccoef(i, k + 1) + &
                Kcoef(i, k + 1) * gama(i, k + 1) - Kcoef(i, k) * gama(i, k)) / buf
        Dcoef(i, k) = Kcoef(i, k) / buf
      END DO
    END DO

    !****************************************************************************************
    ! Niveau k=1

    !****************************************************************************************

    DO i = 1, knon
      buf = delp(i, 1) + Kcoef(i, 2) * (1. - Dcoef(i, 2))
      Acoef(i) = (X(i, 1) * delp(i, 1) + Kcoef(i, 2) * (gama(i, 2) + Ccoef(i, 2))) / buf
      Bcoef(i) = -1. * RG / buf
    END DO

  END SUBROUTINE calc_coef_qbs

  !****************************************************************************************

  SUBROUTINE climb_qbs_up(knon, dtime, qbs_old, &
          flx_qbs1, paprs, pplay, &
          Acoef_QBS_in, Bcoef_QBS_in, &
          Ccoef_QBS_in, Dcoef_QBS_in, &
          Kcoef_qbs_in, gama_qbs_in, &
          flux_qbs, d_qbs)

    ! This routine calculates the flux and tendency of the specific content of blowing snow qbs
    ! The quantity qbs is calculated according to
    ! X(k) = C(k) + D(k)*X(k-1) for X=[qbs], where the coefficients
    ! C and D are known from before and k is index of the vertical layer.

    ! Input arguments
    !****************************************************************************************
    USE lmdz_compbl, ONLY: iflag_pbl, iflag_pbl_split, iflag_order2_sollw, ifl_pbltree
    USE lmdz_yomcst

    IMPLICIT NONE

    INTEGER, INTENT(IN) :: knon
    REAL, INTENT(IN) :: dtime
    REAL, DIMENSION(klon, klev), INTENT(IN) :: qbs_old
    REAL, DIMENSION(klon), INTENT(IN) :: flx_qbs1
    REAL, DIMENSION(klon, klev + 1), INTENT(IN) :: paprs
    REAL, DIMENSION(klon, klev), INTENT(IN) :: pplay

    !!! nrlmd le 02/05/2011
    REAL, DIMENSION(klon), INTENT(IN) :: Acoef_QBS_in, Bcoef_QBS_in
    REAL, DIMENSION(klon, klev), INTENT(IN) :: Ccoef_QBS_in, Dcoef_QBS_in
    REAL, DIMENSION(klon, klev), INTENT(IN) :: Kcoef_qbs_in, gama_qbs_in
    !!!

    ! Output arguments
    !****************************************************************************************
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: flux_qbs, d_qbs

    ! Local variables
    !****************************************************************************************
    LOGICAL, SAVE :: last = .FALSE.
    REAL, DIMENSION(klon, klev) :: qbs_new
    REAL, DIMENSION(klon) :: psref
    INTEGER :: k, i, ierr

    !****************************************************************************************
    ! 1)
    ! Definition of some variables
    REAL, DIMENSION(klon, klev) :: zairm

    !****************************************************************************************
    flux_qbs(:, :) = 0.0
    d_qbs(:, :) = 0.0

    psref(1:knon) = paprs(1:knon, 1)

    IF (mod(iflag_pbl_split, 10) >=1) THEN
      DO i = 1, knon
        Acoef_QBS(i) = Acoef_QBS_in(i)
        Bcoef_QBS(i) = Bcoef_QBS_in(i)
      ENDDO
      DO k = 1, klev
        DO i = 1, knon
          Ccoef_QBS(i, k) = Ccoef_QBS_in(i, k)
          Dcoef_QBS(i, k) = Dcoef_QBS_in(i, k)
          Kcoefqbs(i, k) = Kcoef_qbs_in(i, k)
          IF (k>1) THEN
            gamaqbs(i, k) = gama_qbs_in(i, k)
          ENDIF
        ENDDO
      ENDDO
      !!!
    ENDIF  ! (mod(iflag_pbl_split,2) .ge.1)
    !!!

    !****************************************************************************************
    ! 2)
    ! Calculation of QBS

    !****************************************************************************************

    !- First layer
    qbs_new(1:knon, 1) = Acoef_QBS(1:knon) + Bcoef_QBS(1:knon) * flx_qbs1(1:knon) * dtime
    !- All the other layers
    DO k = 2, klev
      DO i = 1, knon
        qbs_new(i, k) = Ccoef_QBS(i, k) + Dcoef_QBS(i, k) * qbs_new(i, k - 1)
      END DO
    END DO
    !****************************************************************************************
    ! 3)
    ! Calculation of the flux for QBS

    !****************************************************************************************

    !- The flux at first layer, k=1
    flux_qbs(1:knon, 1) = flx_qbs1(1:knon)

    !- The flux at all layers above surface
    DO k = 2, klev
      DO i = 1, knon
        flux_qbs(i, k) = (Kcoefqbs(i, k) / RG / dtime) * &
                (qbs_new(i, k) - qbs_new(i, k - 1) + gamaqbs(i, k))
      END DO
    END DO

    !****************************************************************************************
    ! 4)
    ! Calculation of tendency for QBS

    !****************************************************************************************
    DO k = 1, klev
      DO i = 1, knon
        d_qbs(i, k) = qbs_new(i, k) - qbs_old(i, k)
        zairm(i, k) = (paprs(i, k) - paprs(i, k + 1)) / rg
      END DO
    END DO

    !****************************************************************************************
    ! Some deallocations

    !****************************************************************************************
    IF (last) THEN
      DEALLOCATE(Ccoef_QBS, Dcoef_QBS, stat = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in dealllocate Ccoef_QBS, Dcoef_QBS, ierr=', ierr
      DEALLOCATE(Acoef_QBS, Bcoef_QBS, stat = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in dealllocate Acoef_QBS, Bcoef_QBS, ierr=', ierr
      DEALLOCATE(gamaqbs, stat = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in dealllocate gamaqbs, ierr=', ierr
      DEALLOCATE(Kcoefqbs, stat = ierr)
      IF (ierr /= 0)  PRINT*, ' pb in dealllocate Kcoefqbs, ierr=', ierr
    END IF
  END SUBROUTINE climb_qbs_up

  !****************************************************************************************

END MODULE climb_qbs_mod