!
! $Header$
!
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)
  LOGICAL                            :: firstcall=.TRUE.
  !$OMP THREADPRIVATE(firstcall)

  
  PUBLIC :: climb_wind_down, calcul_wind_flux, 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_gcm(modname,'Pb in allocate alf2',1)

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

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

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

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

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

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

    firstcall=.FALSE.

  END SUBROUTINE climb_wind_init
!
!****************************************************************************************
!
  SUBROUTINE climb_wind_down(knon, dtime, coef_in, pplay, paprs, temp, delp, u_old, v_old)
!
! 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.
!
!
    INCLUDE "YOMCST.h"
! Input arguments
!****************************************************************************************
    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

! 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(:)

! - Calculte the wind components for the first layer
    u1lay(1:knon) = u_old(1:knon,1)*alf1(1:knon) + u_old(1:knon,2)*alf2(1:knon)
    v1lay(1:knon) = v_old(1:knon,1)*alf1(1:knon) + v_old(1:knon,2)*alf2(1:knon)    

! - Calculate K
    Kcoefm(:,:) = 0.0
    DO i = 1, knon
       Kcoefm(i,1) = coef_in(i,1) * (1.0+SQRT(u1lay(i)**2+v1lay(i)**2))* &
            pplay(i,1)/(RD*temp(i,1)) 
       Kcoefm(i,1) = Kcoefm(i,1) * dtime*RG
    END DO

    DO k = 2, klev
       DO i=1,knon
          Kcoefm(i,k) = coef_in(i,k)*RG/(pplay(i,k-1)-pplay(i,k)) &
               *(paprs(i,k)*2/(temp(i,k)+temp(i,k-1))/RD)**2
          Kcoefm(i,k) = Kcoefm(i,k) * dtime*RG
       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(:,:))

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

  END SUBROUTINE climb_wind_down
!
!****************************************************************************************
!
  SUBROUTINE calc_coef(knon, Kcoef, dels, X, alfa1, alfa2, Ccoef, Dcoef)
!
! Find the coefficients C and D in fonction of alfa, K and dels
!
! Input arguments
!****************************************************************************************
    INTEGER, INTENT(IN)                      :: knon
    REAL, DIMENSION(klon,klev), INTENT(IN)   :: Kcoef, dels
    REAL, DIMENSION(klon,klev), INTENT(IN)   :: X
    REAL, DIMENSION(klon), INTENT(IN)        :: alfa1, alfa2

! Output arguments
!****************************************************************************************
    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 = dels(i,klev) + Kcoef(i,klev)

       Ccoef(i,klev) = X(i,klev)*dels(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 = dels(i,k) + Kcoef(i,k) + &
               Kcoef(i,k+1)*(1.-Dcoef(i,k+1))
          
          Ccoef(i,k) = X(i,k)*dels(i,k) + &
               Kcoef(i,k+1)*Ccoef(i,k+1)
          Ccoef(i,k) = Ccoef(i,k)/buf         
          Dcoef(i,k) = Kcoef(i,k)/buf
       END DO
    END DO

! 
! Niveau k=1
!
    DO i = 1, knon
       buf = dels(i,1) + &
            (alfa1(i) + alfa2(i)*Dcoef(i,2)) * Kcoef(i,1) + &
            (1.-Dcoef(i,2))*Kcoef(i,2)
       
       Ccoef(i,1) = X(i,1)*dels(i,1) + &
            (Kcoef(i,2)-Kcoef(i,1)*alfa2(i)) * Ccoef(i,2)
       Ccoef(i,1) = Ccoef(i,1)/buf
       Dcoef(i,1) = Kcoef(i,1)/buf
    END DO

  END SUBROUTINE calc_coef
!
!****************************************************************************************
!
  SUBROUTINE calcul_wind_flux(knon, dtime, &
       flux_u, flux_v)

    INCLUDE "YOMCST.h"

! Input arguments
!****************************************************************************************
    INTEGER, INTENT(IN)                  :: knon
    REAL, INTENT(IN)                     :: dtime

! Output arguments
!****************************************************************************************
    REAL, DIMENSION(klon), INTENT(OUT)   :: flux_u
    REAL, DIMENSION(klon), INTENT(OUT)   :: flux_v

! Local variables
!****************************************************************************************
    INTEGER                              :: i
    REAL, DIMENSION(klon)                :: u0, v0  ! u and v at niveau 0
    REAL, DIMENSION(klon)                :: u1, v1  ! u and v at niveau 1
    REAL, DIMENSION(klon)                :: u2, v2  ! u and v at niveau 2
    

!****************************************************************************************
! Les vents de surface sont supposes nuls
!
!****************************************************************************************
    u0(:) = 0.0
    v0(:) = 0.0

!****************************************************************************************
! On calcule les vents du couhes 1 et 2 recurviement
!
!****************************************************************************************
    DO i = 1, knon
       u1(i) = Ccoef_U(i,1) + Dcoef_U(i,1)*u0(i)
       v1(i) = Ccoef_V(i,1) + Dcoef_V(i,1)*v0(i)
       u2(i) = Ccoef_U(i,2) + Dcoef_U(i,2)*u1(i)
       v2(i) = Ccoef_V(i,2) + Dcoef_V(i,2)*v1(i)
    END DO

!****************************************************************************************
! On calcule le flux
!
!****************************************************************************************
    flux_u(:) = 0.0
    flux_v(:) = 0.0

    DO i=1,knon
       flux_u(i) = Kcoefm(i,1)/RG/dtime * (u1(i)*alf1(i) + u2(i)*alf2(i) - u0(i)) 
       flux_v(i) = Kcoefm(i,1)/RG/dtime * (v1(i)*alf1(i) + v2(i)*alf2(i) - v0(i)) 
    END DO

  END SUBROUTINE calcul_wind_flux
!
!****************************************************************************************
!
  SUBROUTINE climb_wind_up(knon, dtime, u_old, v_old, &
       flx_u_new, flx_v_new, d_u_new, d_v_new)
!
! Diffuse the wind components from the surface and up to the top layer. Coefficents
! C and D are known from before. The values for U and V at surface are supposed to be 
! zero (this could be modified).
!
! u(k) = Cu(k) + Du(k)*u(k-1)
! v(k) = Cv(k) + Dv(k)*v(k-1)
! [1 <= k <= klev]
!
!****************************************************************************************
    INCLUDE "YOMCST.h"

! Input arguments
!****************************************************************************************
    INTEGER, INTENT(IN)                     :: knon
    REAL, INTENT(IN)                        :: dtime
    REAL, DIMENSION(klon,klev), INTENT(IN)  :: u_old
    REAL, DIMENSION(klon,klev), INTENT(IN)  :: v_old

! 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
    REAL, DIMENSION(klon)                   :: u0, v0
    INTEGER                                 :: k, i
    
!
!****************************************************************************************

! Niveau 0
    u0(1:knon) = 0.0
    v0(1:knon) = 0.0

! Niveau 1
    DO i = 1, knon
       u_new(i,1) = Ccoef_U(i,1) + Dcoef_U(i,1) * u0(i)
       v_new(i,1) = Ccoef_V(i,1) + Dcoef_V(i,1) * v0(i)
    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

! Niveau 1
    DO i = 1, knon
       flx_u_new(i,1) = Kcoefm(i,1)/RG/dtime * &
            (u_new(i,1)*alf1(i)+u_new(i,2)*alf2(i) - u0(i)) 
       flx_v_new(i,1) = Kcoefm(i,1)/RG/dtime * &
            (v_new(i,1)*alf1(i)+v_new(i,2)*alf2(i) - v0(i)) 
    END DO

! 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