MODULE climb_wind_mod

  ! Module to solve the verctical diffusion of the wind components "u" and "v".

  USE dimphy
  USE lmdz_abort_physic, ONLY: abort_physic

  IMPLICIT NONE

  SAVE
  PRIVATE

  REAL, DIMENSION(:), ALLOCATABLE :: alf1, alf2
  !$OMP THREADPRIVATE(alf1,alf2)
  REAL, DIMENSION(:, :), ALLOCATABLE :: Kcoefm
  !$OMP THREADPRIVATE(Kcoefm)
  REAL, DIMENSION(:, :), ALLOCATABLE :: Ccoef_U, Dcoef_U
  !$OMP THREADPRIVATE(Ccoef_U, Dcoef_U)
  REAL, DIMENSION(:, :), ALLOCATABLE :: Ccoef_V, Dcoef_V
  !$OMP THREADPRIVATE(Ccoef_V, Dcoef_V)
  REAL, DIMENSION(:), ALLOCATABLE :: Acoef_U, Bcoef_U
  !$OMP THREADPRIVATE(Acoef_U, Bcoef_U)
  REAL, DIMENSION(:), ALLOCATABLE :: Acoef_V, Bcoef_V
  !$OMP THREADPRIVATE(Acoef_V, Bcoef_V)
  LOGICAL :: firstcall = .TRUE.
  !$OMP THREADPRIVATE(firstcall)

  PUBLIC :: climb_wind_down, climb_wind_up

CONTAINS

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

  SUBROUTINE climb_wind_init

    INTEGER :: ierr
    CHARACTER(len = 20) :: modname = 'climb_wind_init'

    !****************************************************************************************
    ! Allocation of global module variables

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

    ALLOCATE(alf1(klon), stat = ierr)
    IF (ierr /= 0) CALL abort_physic(modname, 'Pb in allocate alf1', 1)

    ALLOCATE(alf2(klon), stat = ierr)
    IF (ierr /= 0) CALL abort_physic(modname, 'Pb in allocate alf2', 1)

    ALLOCATE(Kcoefm(klon, klev), stat = ierr)
    IF (ierr /= 0) CALL abort_physic(modname, 'Pb in allocate Kcoefm', 1)

    ALLOCATE(Ccoef_U(klon, klev), stat = ierr)
    IF (ierr /= 0) CALL abort_physic(modname, 'Pb in allocate Ccoef_U', 1)

    ALLOCATE(Dcoef_U(klon, klev), stat = ierr)
    IF (ierr /= 0) CALL abort_physic(modname, 'Pb in allocation Dcoef_U', 1)

    ALLOCATE(Ccoef_V(klon, klev), stat = ierr)
    IF (ierr /= 0) CALL abort_physic(modname, 'Pb in allocation Ccoef_V', 1)

    ALLOCATE(Dcoef_V(klon, klev), stat = ierr)
    IF (ierr /= 0) CALL abort_physic(modname, 'Pb in allocation Dcoef_V', 1)

    ALLOCATE(Acoef_U(klon), Bcoef_U(klon), Acoef_V(klon), Bcoef_V(klon), STAT = ierr)
    IF (ierr /= 0)  PRINT*, ' pb in allloc Acoef_U and Bcoef_U, ierr=', ierr

    firstcall = .FALSE.

  END SUBROUTINE climb_wind_init

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

  SUBROUTINE climb_wind_down(knon, dtime, coef_in, pplay, paprs, temp, delp, u_old, v_old, &
          !!! nrlmd le 02/05/2011
          Ccoef_U_out, Ccoef_V_out, Dcoef_U_out, Dcoef_V_out, &
          Kcoef_m_out, alf_1_out, alf_2_out, &
          !!!
          Acoef_U_out, Acoef_V_out, Bcoef_U_out, Bcoef_V_out)

    ! This routine calculates for the wind components u and v,
    ! recursivly the coefficients C and D in equation
    ! X(k) = C(k) + D(k)*X(k-1), X=[u,v], k=[1,klev] is 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) :: coef_in
    REAL, DIMENSION(klon, klev), INTENT(IN) :: pplay ! pres au milieu de couche (Pa)
    REAL, DIMENSION(klon, klev + 1), INTENT(IN) :: paprs ! pression a inter-couche (Pa)
    REAL, DIMENSION(klon, klev), INTENT(IN) :: temp  ! temperature
    REAL, DIMENSION(klon, klev), INTENT(IN) :: delp
    REAL, DIMENSION(klon, klev), INTENT(IN) :: u_old
    REAL, DIMENSION(klon, klev), INTENT(IN) :: v_old

    ! Output arguments
    !****************************************************************************************
    REAL, DIMENSION(klon), INTENT(OUT) :: Acoef_U_out
    REAL, DIMENSION(klon), INTENT(OUT) :: Acoef_V_out
    REAL, DIMENSION(klon), INTENT(OUT) :: Bcoef_U_out
    REAL, DIMENSION(klon), INTENT(OUT) :: Bcoef_V_out

    !!! nrlmd le 02/05/2011
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Ccoef_U_out
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Ccoef_V_out
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Dcoef_U_out
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Dcoef_V_out
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: Kcoef_m_out
    REAL, DIMENSION(klon), INTENT(OUT) :: alf_1_out
    REAL, DIMENSION(klon), INTENT(OUT) :: alf_2_out
    !!!

    ! Local variables
    !****************************************************************************************
    REAL, DIMENSION(klon) :: u1lay, v1lay
    INTEGER :: k, i

    !****************************************************************************************
    ! Initialize module
    IF (firstcall) CALL climb_wind_init

    !****************************************************************************************
    ! Calculate the coefficients C and D in : u(k) = C(k) + D(k)*u(k-1)

    !****************************************************************************************
    ! - Define alpha (alf1 and alf2)
    alf1(:) = 1.0
    alf2(:) = 1.0 - alf1(:)

    ! - Calculate the coefficients K
    Kcoefm(:, :) = 0.0
    DO k = 2, klev
      DO i = 1, knon
        Kcoefm(i, k) = coef_in(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
      END DO
    END DO

    ! - Calculate the coefficients C and D, component "u"
    CALL calc_coef(knon, Kcoefm(:, :), delp(:, :), &
            u_old(:, :), alf1(:), alf2(:), &
            Ccoef_U(:, :), Dcoef_U(:, :), Acoef_U(:), Bcoef_U(:))

    ! - Calculate the coefficients C and D, component "v"
    CALL calc_coef(knon, Kcoefm(:, :), delp(:, :), &
            v_old(:, :), alf1(:), alf2(:), &
            Ccoef_V(:, :), Dcoef_V(:, :), Acoef_V(:), Bcoef_V(:))

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

    !****************************************************************************************
    Acoef_U_out = Acoef_U
    Bcoef_U_out = Bcoef_U
    Acoef_V_out = Acoef_V
    Bcoef_V_out = Bcoef_V

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

    !****************************************************************************************
    !!! jyg le 07/02/2012
    !!jyg       IF (mod(iflag_pbl_split,2) .EQ.1) THEN
    IF (mod(iflag_pbl_split, 10) >=1) THEN
      !!! nrlmd le 02/05/2011
      DO k = 1, klev
        DO i = 1, klon
          Ccoef_U_out(i, k) = Ccoef_U(i, k)
          Ccoef_V_out(i, k) = Ccoef_V(i, k)
          Dcoef_U_out(i, k) = Dcoef_U(i, k)
          Dcoef_V_out(i, k) = Dcoef_V(i, k)
          Kcoef_m_out(i, k) = Kcoefm(i, k)
        ENDDO
      ENDDO
      DO i = 1, klon
        alf_1_out(i) = alf1(i)
        alf_2_out(i) = alf2(i)
      ENDDO
      !!!
    ENDIF  ! (mod(iflag_pbl_split,2) .ge.1)
    !!!

  END SUBROUTINE climb_wind_down

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

  SUBROUTINE calc_coef(knon, Kcoef, delp, X, alfa1, alfa2, Ccoef, Dcoef, Acoef, Bcoef)
    USE lmdz_yomcst

    IMPLICIT NONE

    ! Find the coefficients C and D in fonction of alfa, K and delp

    ! Input arguments
    !****************************************************************************************
    INTEGER, INTENT(IN) :: knon
    REAL, DIMENSION(klon, klev), INTENT(IN) :: Kcoef, delp
    REAL, DIMENSION(klon, klev), INTENT(IN) :: X
    REAL, DIMENSION(klon), INTENT(IN) :: alfa1, alfa2

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

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

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

    ! Calculate coefficients C and D at top level, 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) / buf
      Dcoef(i, klev) = Kcoef(i, klev) / buf
    END DO

    ! Calculate coefficients C and D at top level (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)) / buf
        Dcoef(i, k) = Kcoef(i, k) / buf
      END DO
    END DO

    ! Calculate coeffiecent A and B at surface

    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) * Ccoef(i, 2)) / buf
      Bcoef(i) = -RG / buf
    END DO

  END SUBROUTINE calc_coef

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

  SUBROUTINE climb_wind_up(knon, dtime, u_old, v_old, flx_u1, flx_v1, &
          !!! nrlmd le 02/05/2011
          Acoef_U_in, Acoef_V_in, Bcoef_U_in, Bcoef_V_in, &
          Ccoef_U_in, Ccoef_V_in, Dcoef_U_in, Dcoef_V_in, &
          Kcoef_m_in, &
          !!!
          flx_u_new, flx_v_new, d_u_new, d_v_new)

    ! Diffuse the wind components from the surface layer and up to the top layer.
    ! Coefficents A, B, C and D are known from before. Start values for the diffusion are the
    ! momentum fluxes at surface.

    ! u(k=1) = A + B*flx*dtime
    ! u(k)   = C(k) + D(k)*u(k-1)  [2 <= k <= klev]

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

    ! 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) :: u_old
    REAL, DIMENSION(klon, klev), INTENT(IN) :: v_old
    REAL, DIMENSION(klon), INTENT(IN) :: flx_u1, flx_v1 ! momentum flux

    !!! nrlmd le 02/05/2011
    REAL, DIMENSION(klon), INTENT(IN) :: Acoef_U_in, Acoef_V_in, Bcoef_U_in, Bcoef_V_in
    REAL, DIMENSION(klon, klev), INTENT(IN) :: Ccoef_U_in, Ccoef_V_in, Dcoef_U_in, Dcoef_V_in
    REAL, DIMENSION(klon, klev), INTENT(IN) :: Kcoef_m_in
    !!!

    ! Output arguments
    !****************************************************************************************
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: flx_u_new, flx_v_new
    REAL, DIMENSION(klon, klev), INTENT(OUT) :: d_u_new, d_v_new

    ! Local variables
    !****************************************************************************************
    REAL, DIMENSION(klon, klev) :: u_new, v_new
    INTEGER :: k, i
    !****************************************************************************************

    !!! jyg le 07/02/2012
    !!jyg       IF (mod(iflag_pbl_split,2) .EQ.1) THEN
    IF (mod(iflag_pbl_split, 10) >=1) THEN
      !!! nrlmd le 02/05/2011
      DO i = 1, knon
        Acoef_U(i) = Acoef_U_in(i)
        Acoef_V(i) = Acoef_V_in(i)
        Bcoef_U(i) = Bcoef_U_in(i)
        Bcoef_V(i) = Bcoef_V_in(i)
      ENDDO
      DO k = 1, klev
        DO i = 1, knon
          Ccoef_U(i, k) = Ccoef_U_in(i, k)
          Ccoef_V(i, k) = Ccoef_V_in(i, k)
          Dcoef_U(i, k) = Dcoef_U_in(i, k)
          Dcoef_V(i, k) = Dcoef_V_in(i, k)
          Kcoefm(i, k) = Kcoef_m_in(i, k)
        ENDDO
      ENDDO
      !!!
    ENDIF  ! (mod(iflag_pbl_split,2) .ge.1)
    !!!

    ! Niveau 1
    DO i = 1, knon
      u_new(i, 1) = Acoef_U(i) + Bcoef_U(i) * flx_u1(i) * dtime
      v_new(i, 1) = Acoef_V(i) + Bcoef_V(i) * flx_v1(i) * dtime
    END DO

    ! Niveau 2 jusqu'au sommet klev
    DO k = 2, klev
      DO i = 1, knon
        u_new(i, k) = Ccoef_U(i, k) + Dcoef_U(i, k) * u_new(i, k - 1)
        v_new(i, k) = Ccoef_V(i, k) + Dcoef_V(i, k) * v_new(i, k - 1)
      END DO
    END DO

    !****************************************************************************************
    ! Calcul flux

    !== flux_u/v est le flux de moment angulaire (positif vers bas)
    !== dont l'unite est: (kg m/s)/(m**2 s)

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

    flx_u_new(:, :) = 0.0
    flx_v_new(:, :) = 0.0

    flx_u_new(1:knon, 1) = flx_u1(1:knon)
    flx_v_new(1:knon, 1) = flx_v1(1:knon)

    ! Niveau 2->klev
    DO k = 2, klev
      DO i = 1, knon
        flx_u_new(i, k) = Kcoefm(i, k) / RG / dtime * &
                (u_new(i, k) - u_new(i, k - 1))

        flx_v_new(i, k) = Kcoefm(i, k) / RG / dtime * &
                (v_new(i, k) - v_new(i, k - 1))
      END DO
    END DO

    !****************************************************************************************
    ! Calcul tendances

    !****************************************************************************************
    d_u_new(:, :) = 0.0
    d_v_new(:, :) = 0.0
    DO k = 1, klev
      DO i = 1, knon
        d_u_new(i, k) = u_new(i, k) - u_old(i, k)
        d_v_new(i, k) = v_new(i, k) - v_old(i, k)
      END DO
    END DO

  END SUBROUTINE climb_wind_up

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

END MODULE climb_wind_mod