!
MODULE climb_wind_mod
!
! Module to solve the verctical diffusion of the wind components "u" and "v".
!
  USE dimphy

  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, climb_wind_init

CONTAINS
!
!****************************************************************************************
!
  SUBROUTINE climb_wind_init

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


!****************************************************************************************
! Initialize module
    
    IF (firstcall) THEN

      !****************************************************************************************
      ! Allocation of global module variables
      !
      !****************************************************************************************
  
      ALLOCATE(alf1(klon), stat=ierr)
      IF (ierr /= 0) CALL abort_physic(modname,'Pb in allocate alf1',1)
      alf1(:) = 0.
  
      ALLOCATE(alf2(klon), stat=ierr)
      IF (ierr /= 0) CALL abort_physic(modname,'Pb in allocate alf2',1)
      alf2(:) = 0.
  
      ALLOCATE(Kcoefm(klon,klev), stat=ierr)
      IF (ierr /= 0) CALL abort_physic(modname,'Pb in allocate Kcoefm',1)
      Kcoefm(:,:) = 0.
  
      ALLOCATE(Ccoef_U(klon,klev), stat=ierr)
      IF (ierr /= 0) CALL abort_physic(modname,'Pb in allocate Ccoef_U',1)
      Ccoef_U(:,:) = 0.
  
      ALLOCATE(Dcoef_U(klon,klev), stat=ierr)
      IF (ierr /= 0) CALL abort_physic(modname,'Pb in allocation Dcoef_U',1)
      Dcoef_U(:,:) = 0.
  
      ALLOCATE(Ccoef_V(klon,klev), stat=ierr)
      IF (ierr /= 0) CALL abort_physic(modname,'Pb in allocation Ccoef_V',1)
      Ccoef_V(:,:) = 0.
  
      ALLOCATE(Dcoef_V(klon,klev), stat=ierr)
      IF (ierr /= 0) CALL abort_physic(modname,'Pb in allocation Dcoef_V',1)
      Dcoef_V(:,:) = 0.
  
      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
      Acoef_U(:) = 0. ; Bcoef_U(:) = 0. ; Acoef_V(:) = 0. ; Bcoef_V(:) = 0. ;
  
      firstcall=.FALSE.
    
    ENDIF

  END SUBROUTINE climb_wind_init
!
!****************************************************************************************
!
  SUBROUTINE climb_wind_down(knon, ni, 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)
!$gpum horizontal knon

!
! 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.
!
!
USE yomcst_mod_h
USE compbl_mod_h
! Input arguments
!****************************************************************************************
    INTEGER, INTENT(IN)                      :: knon
    INTEGER, INTENT(IN)                      :: ni(knon)
    REAL, INTENT(IN)                         :: dtime
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: coef_in
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: pplay ! pres au milieu de couche (Pa)
    REAL, DIMENSION(knon,klev+1), INTENT(IN) :: paprs ! pression a inter-couche (Pa)
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: temp  ! temperature
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: delp
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: u_old
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: v_old

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

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

! Local variables
!****************************************************************************************
    REAL :: yalf1(knon)
    REAL :: yalf2(knon)
    REAL :: yKcoefm(knon,klev)
    REAL :: yCcoef_U(knon,klev)
    REAL :: yDcoef_U(knon,klev)
    REAL :: yCcoef_V(knon,klev)
    REAL :: yDcoef_V(knon,klev)
    REAL :: yAcoef_U(knon), yBcoef_U(knon), yAcoef_V(knon), yBcoef_V(knon)

    REAL, DIMENSION(knon)                    :: u1lay, v1lay
    INTEGER                                  :: k, i, j


!****************************************************************************************
! Calculate the coefficients C and D in : u(k) = C(k) + D(k)*u(k-1)
!
!****************************************************************************************
! - Define alpha (alf1 and alf2) 
    yalf1(:) = 1.0
    yalf2(:) = 1.0 - yalf1(:)

! - Calculate the coefficients K
    yKcoefm(:,:) = 0.0
    DO k = 2, klev
       DO i=1,knon
          yKcoefm(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, yKcoefm, delp, &
         u_old, yalf1, yalf2,  &
         yCcoef_U, yDcoef_U, yAcoef_U, yBcoef_U)

! - Calculate the coefficients C and D, component "v"
    CALL calc_coef(knon, yKcoefm, delp, &
         v_old, yalf1, yalf2,  &
         yCcoef_V, yDcoef_V, yAcoef_V, yBcoef_V)

!****************************************************************************************
! 6)
! Return the first layer in output variables
!
!****************************************************************************************
    Acoef_U_out = yAcoef_U
    Bcoef_U_out = yBcoef_U
    Acoef_V_out = yAcoef_V
    Bcoef_V_out = yBcoef_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) .ge.1) THEN
!!! nrlmd le 02/05/2011
    DO k= 1, klev
      DO i= 1, knon
        Ccoef_U_out(i,k) = yCcoef_U(i,k)
        Ccoef_V_out(i,k) = yCcoef_V(i,k)
        Dcoef_U_out(i,k) = yDcoef_U(i,k)
        Dcoef_V_out(i,k) = yDcoef_V(i,k)
        Kcoef_m_out(i,k) = yKcoefm(i,k)
      ENDDO
    ENDDO
    DO i= 1, knon
      alf_1_out(i)   = yalf1(i)
      alf_2_out(i)   = yalf2(i)
    ENDDO
!!!      
       ENDIF  ! (mod(iflag_pbl_split,2) .ge.1)
!!!
    DO j= 1, knon
      i=ni(j)
      Acoef_U(i) = yAcoef_U(j)
      Bcoef_U(i) = yBcoef_U(j)
      Acoef_V(i) = yAcoef_V(j)
      Bcoef_V(i) = yBcoef_V(j)
    ENDDO

    DO k= 1, klev
      DO j= 1, knon
        i=ni(j)
        Ccoef_U(i,k) = yCcoef_U(j,k)
        Ccoef_V(i,k) = yCcoef_V(j,k)
        Dcoef_U(i,k) = yDcoef_U(j,k)
        Dcoef_V(i,k) = yDcoef_V(j,k)
        Kcoefm(i,k) = yKcoefm(j,k)
        alf1(i)   = yalf1(j)
        alf2(i)   = yalf2(j)
      ENDDO
    ENDDO
  END SUBROUTINE climb_wind_down
!
!****************************************************************************************
!
  SUBROUTINE calc_coef(knon, Kcoef, delp, X, alfa1, alfa2, Ccoef, Dcoef, Acoef, Bcoef)
!$gpum horizontal knon

!
! Find the coefficients C and D in fonction of alfa, K and delp
USE yomcst_mod_h
! Input arguments
!****************************************************************************************
    INTEGER, INTENT(IN)                      :: knon
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: Kcoef, delp
    REAL, DIMENSION(knon,klev), INTENT(IN)   :: X
    REAL, DIMENSION(knon), INTENT(IN)        :: alfa1, alfa2

! Output arguments
!****************************************************************************************
    REAL, DIMENSION(knon), INTENT(OUT)       :: Acoef, Bcoef
    REAL, DIMENSION(knon,klev), INTENT(OUT)  :: Ccoef, Dcoef
  
! local variables
!****************************************************************************************
    INTEGER                                  :: k, i, j
    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, ni, 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)
!$gpum horizontal knon

!
! 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]
!
!****************************************************************************************
USE yomcst_mod_h
USE compbl_mod_h
! Input arguments
!****************************************************************************************
    INTEGER, INTENT(IN)                     :: knon
    INTEGER, INTENT(IN)                     :: ni(knon)
    REAL, INTENT(IN)                        :: dtime
    REAL, DIMENSION(knon,klev), INTENT(IN)  :: u_old
    REAL, DIMENSION(knon,klev), INTENT(IN)  :: v_old
    REAL, DIMENSION(knon), INTENT(IN)       :: flx_u1, flx_v1 ! momentum flux

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

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

! Local variables
!****************************************************************************************
    REAL :: yalf1(knon)
    REAL :: yalf2(knon)
    REAL :: yKcoefm(knon,klev)
    REAL :: yCcoef_U(knon,klev)
    REAL :: yDcoef_U(knon,klev)
    REAL :: yCcoef_V(knon,klev)
    REAL :: yDcoef_V(knon,klev)
    REAL :: yAcoef_U(knon), yBcoef_U(knon), yAcoef_V(knon), yBcoef_V(knon)

    REAL, DIMENSION(knon,klev)              :: u_new, v_new
    INTEGER                                 :: k, i, j
!****************************************************************************************

!!! jyg le 07/02/2012
!!jyg       IF (mod(iflag_pbl_split,2) .eq.1) THEN
       IF (mod(iflag_pbl_split,10) .ge.1) THEN
!!! nrlmd le 02/05/2011
    DO i = 1, knon
      yAcoef_U(i)=Acoef_U_in(i)
      yAcoef_V(i)=Acoef_V_in(i)
      yBcoef_U(i)=Bcoef_U_in(i)
      yBcoef_V(i)=Bcoef_V_in(i)
    ENDDO
    DO k = 1, klev
      DO i = 1, knon
        yCcoef_U(i,k)=Ccoef_U_in(i,k)
        yCcoef_V(i,k)=Ccoef_V_in(i,k)
        yDcoef_U(i,k)=Dcoef_U_in(i,k)
        yDcoef_V(i,k)=Dcoef_V_in(i,k)
        yKcoefm(i,k)=Kcoef_m_in(i,k)
      ENDDO
    ENDDO
   ELSE
    DO j = 1, knon
      i=ni(j)
      yAcoef_U(j)=Acoef_U(i)
      yAcoef_V(j)=Acoef_V(i)
      yBcoef_U(j)=Bcoef_U(i)
      yBcoef_V(j)=Bcoef_V(i)
    ENDDO
    DO k = 1, klev
      DO j = 1, knon
        i=ni(j)
        yCcoef_U(j,k)=Ccoef_U(i,k)
        yCcoef_V(j,k)=Ccoef_V(i,k)
        yDcoef_U(j,k)=Dcoef_U(i,k)
        yDcoef_V(j,k)=Dcoef_V(i,k)
        yKcoefm(j,k)=Kcoefm(i,k)
      ENDDO
    ENDDO    
!!!
       ENDIF  ! (mod(iflag_pbl_split,2) .ge.1)
!!!

! Niveau 1
    DO i = 1, knon
       u_new(i,1) = yAcoef_U(i) + yBcoef_U(i)*flx_u1(i)*dtime
       v_new(i,1) = yAcoef_V(i) + yBcoef_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) = yCcoef_U(i,k) + yDcoef_U(i,k) * u_new(i,k-1)
          v_new(i,k) = yCcoef_V(i,k) + yDcoef_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) = yKcoefm(i,k)/RG/dtime * &
               (u_new(i,k)-u_new(i,k-1))
          
          flx_v_new(i,k) = yKcoefm(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

    DO j = 1, knon
      i=ni(j)
      Acoef_U(i)=yAcoef_U(j)
      Acoef_V(i)=yAcoef_V(j)
      Bcoef_U(i)=yBcoef_U(j)
      Bcoef_V(i)=yBcoef_V(j)
    ENDDO

    DO k = 1, klev
      DO j = 1, knon
        i=ni(j)
        Ccoef_U(i,k) = yCcoef_U(j,k)
        Ccoef_V(i,k) = yCcoef_V(j,k)
        Dcoef_U(i,k) = yDcoef_U(j,k)
        Dcoef_V(i,k) = yDcoef_V(j,k)
        Kcoefm(i,k) = yKcoefm(j,k) 
      ENDDO
    ENDDO    
  
  END SUBROUTINE climb_wind_up
!
!****************************************************************************************
!
END MODULE climb_wind_mod