module atke_exchange_coeff_mod

implicit none

contains 

subroutine atke_compute_km_kh(ngrid,nlay, &
                        wind_u,wind_v,temp,play,pinterf, &
                        tke,Km,Kh)

!========================================================================
! Routine that computes turbulent Km / Kh coefficients with a
! 1.5 order closure scheme (TKE) with or without stationarity assumption
!
! This parameterization has been constructed in the framework of a
! collective and collaborative workshop, 
! the so-called 'Atelier TKE (ATKE)' with
! K. Arjdal, L. Raillard, C. Dehondt, P. Tiengou, A. Spiga, F. Cheruy, 
! M. Coulon-Decorzens, S. Fromang, G. Riviere, A. Sima, F. Hourdin, E. Vignon
!
! Main assumptions of the model :
! (1) dry atmosphere
! (2) horizontal homogeneity (Dx=Dy=0.)
!=======================================================================



USE atke_turbulence_ini_mod, ONLY : iflag_atke, kappa, l0, ric, cn, cinf, rpi
USE atke_turbulence_ini_mod, ONLY : cepsilon, pr_slope, pr_asym, rg, rd

implicit none


! Declarations:
!=============

INTEGER, INTENT(IN) :: ngrid ! number of horizontal index (flat grid)
INTEGER, INTENT(IN) :: nlay ! number of vertical index  

REAL, DIMENSION(ngrid,nlay), INTENT(IN)       :: wind_u   ! zonal velocity (m/s)
REAL, DIMENSION(ngrid,nlay), INTENT(IN)       :: wind_v   ! meridional velocity (m/s)
REAL, DIMENSION(ngrid,nlay), INTENT(IN)       :: temp   ! temperature (K)
REAL, DIMENSION(ngrid,nlay), INTENT(IN)       :: play   ! pressure (Pa)
REAL, DIMENSION(ngrid,nlay+1), INTENT(IN)     :: pinterf   ! pressure at interfaces(Pa)


REAL, DIMENSION(ngrid,nlay+1), INTENT(INOUT)  :: tke  ! turbulent kinetic energy at interface between layers

REAL, DIMENSION(ngrid,nlay+1), INTENT(OUT)    :: Km   ! Exchange coefficient for momentum at interface between layers
REAL, DIMENSION(ngrid,nlay+1), INTENT(OUT)    :: Kh   ! Exchange coefficient for heat flux at interface between layers

! Local variables
REAL, DIMENSION(ngrid,nlay+1) :: l_exchange ! Length of exchange
REAL, DIMENSION(ngrid,nlay+1) :: z_interf ! Altitude at the interface
REAL, DIMENSION(ngrid,nlay+1) :: dz_interf ! distance between two consecutive interfaces
REAL, DIMENSION(ngrid,nlay+1) :: Ri ! Richardson's number
REAL, DIMENSION(ngrid,nlay+1) :: Prandtl ! Turbulent Prandtl's number 
REAL, DIMENSION(ngrid,nlay+1) :: Sm ! Stability function for momentum
REAL, DIMENSION(ngrid,nlay+1) :: Sh ! Stability function for heat

INTEGER :: igrid,ilay ! horizontal,vertical index (flat grid)
REAL    :: Ri0        ! parameter for Sm stability function



! Initializations:
!================

DO igrid=1,ngrid
    z_interf(igrid,1) = 0.0
    Ri(igrid,1) = 0.1! TO DO
    Sm(igrid,1) = 0.0 ! TO DO
END DO



! Calculation of altitude of layers' bottom interfaces:
!=====================================================

DO ilay=2,nlay+1
    DO igrid=1,ngrid
        dz_interf(igrid,ilay) = rg * temp(igrid,ilay) / rg * log(play(igrid,ilay-1)/play(igrid,ilay))
        z_interf(igrid,ilay) = z_interf(igrid,ilay-1) + dz_interf(igrid,ilay)
    ENDDO
ENDDO


! Computing the mixing length:
!=============================

DO ilay=2,nlay+1
    DO igrid=1,ngrid
        l_exchange(igrid,ilay) = kappa*l0*z_interf(igrid,ilay) / (kappa*z_interf(igrid,ilay) + l0)
    ENDDO
ENDDO

! Computes the gradient Richardson's number and stability functions:
!===================================================================

DO ilay=2,nlay+1
    DO igrid=1,ngrid
        ! Warning: sure that gradient of temp and wind should be calculated with dz_interf and not with dz_lay?
        Ri(igrid,ilay) = rg * (temp(igrid,ilay) - temp(igrid,ilay-1)) / dz_interf(igrid,ilay) / & 
        (((wind_u(igrid,ilay) - wind_u(igrid,ilay-1)) / dz_interf(igrid,ilay))**2 + & 
        ((wind_v(igrid,ilay) - wind_v(igrid,ilay-1)) / dz_interf(igrid,ilay))**2 ) 

        ! calculation of Ri0 such that continuity in slope of Sm at Ri=0
        Ri0=2.*ric*cinf/rpi/cn

        IF (Ri(igrid,ilay) < 0.) THEN
            Sm(igrid,ilay) = 2./rpi * cinf * atan(-Ri(igrid,ilay)/Ri0) + cn
            Prandtl(igrid,ilay) = 1. + Ri(igrid,ilay) * pr_slope
        ELSE
            Sm(igrid,ilay) = max(0.,cn*(1.-Ri(igrid,ilay)/Ric))
            Prandtl(igrid,ilay) = 2./rpi * atan(-Ri(igrid,ilay)) - 1. + pr_asym  ! exp(Ri(igrid,ilay))
        END IF
        
        Sh(igrid,ilay) = Sm(igrid,ilay) / Prandtl(igrid,ilay)

    ENDDO
ENDDO

! Computing the TKE:
!===================
IF (iflag_atke == 0) THEN

! stationary solution neglecting the vertical transport of TKE by turbulence
    DO ilay=2,nlay+1
        DO igrid=1,ngrid
            tke(igrid,ilay) = cepsilon * l_exchange(igrid,ilay)**2 * Sm(igrid,ilay) * &
            (((wind_u(igrid,ilay) - wind_u(igrid,ilay-1)) / dz_interf(igrid,ilay))**2 + & 
            ((wind_v(igrid,ilay) - wind_v(igrid,ilay-1)) / dz_interf(igrid,ilay))**2 ) * &
            (1. - Ri(igrid,ilay) / Prandtl(igrid,ilay))
        ENDDO
    ENDDO

ELSE ! TO DO
    
   call abort_physic("atke_compute_km_kh", &
        'traitement non-stationnaire de la tke pas encore prevu', 1)

END IF


! Computing eddy diffusivity coefficients:
!========================================
DO ilay=2,nlay+1
    DO igrid=1,ngrid
        Km(igrid,ilay) = l_exchange(igrid,ilay) * Sm(igrid,ilay) * tke(igrid,ilay)**0.5
        Kh(igrid,ilay) = l_exchange(igrid,ilay) * Sh(igrid,ilay) * tke(igrid,ilay)**0.5
    END DO
END DO



end subroutine atke_compute_km_kh


end module atke_exchange_coeff_mod