module lmdz_call_atke

USE lmdz_atke_exchange_coeff, ONLY :  atke_compute_km_kh

implicit none


contains

subroutine call_atke(dtime,ngrid,nlay,cdrag_uv,cdrag_t,u_surf,v_surf,temp_surf, &
                        wind_u,wind_v,temp,qvap,play,pinterf,zlay,zinterf, &
                        tke,Km_out,Kh_out)




USE lmdz_atke_turbulence_ini, ONLY : iflag_num_atke, rg, rd

implicit none


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

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


REAL, DIMENSION(ngrid), INTENT(IN)       :: cdrag_uv  ! drag coefficient for wind
REAL, DIMENSION(ngrid), INTENT(IN)       :: cdrag_t   ! drag coefficient for temperature

REAL, DIMENSION(ngrid), INTENT(IN)       :: u_surf    ! x wind velocity at the surface
REAL, DIMENSION(ngrid), INTENT(IN)       :: v_surf    ! y wind velocity at the surface
REAL, DIMENSION(ngrid), INTENT(IN)       :: temp_surf ! surface temperature

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)       :: qvap   ! specific humidity (kg/kg)
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), INTENT(IN)       :: zlay      ! altitude at the middle of the vertical layers (m)
REAL, DIMENSION(ngrid,nlay+1), INTENT(IN)     :: zinterf   ! altitude at interfaces (m)


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

REAL, DIMENSION(ngrid,nlay), INTENT(OUT)      :: Km_out   ! output: Exchange coefficient for momentum at interface between layers
REAL, DIMENSION(ngrid,nlay), INTENT(OUT)      :: Kh_out   ! output: Exchange coefficient for heat flux at interface between layers


REAL, DIMENSION(ngrid,nlay) :: wind_u_predict, wind_v_predict
REAL, DIMENSION(ngrid) ::  wind1
INTEGER i


call atke_compute_km_kh(ngrid,nlay,dtime,&
                        wind_u,wind_v,temp,qvap,play,pinterf,zlay,zinterf,cdrag_uv,&
                        tke,Km_out,Kh_out)


                               
if (iflag_num_atke .EQ. 1) then
   !! In this case, we make an explicit prediction of the wind shear to calculate the tke in a
   !! forward backward way
   !! pay attention that the treatment of the TKE
   !! has to be adapted when solving the TKE with a prognostic equation

   do i=1,ngrid
      wind1(i)=sqrt(wind_u(i,1)**2+wind_v(i,1)**2)
   enddo
   call atke_explicit_prediction(ngrid,nlay,rg,rd,dtime,pinterf,play,temp,wind1,wind_u,Km_out,u_surf,cdrag_uv,wind_u_predict)  
   call atke_explicit_prediction(ngrid,nlay,rg,rd,dtime,pinterf,play,temp,wind1,wind_v,Km_out,v_surf,cdrag_uv,wind_v_predict)


   call atke_compute_km_kh(ngrid,nlay,dtime,&
                        wind_u_predict,wind_v_predict,temp,qvap,play,pinterf,zlay,zinterf,cdrag_uv, &
                        tke,Km_out,Kh_out)

end if






end subroutine call_atke

!----------------------------------------------------------------------------------------

subroutine atke_explicit_prediction(ngrid,nlay,rg,rd,dtime,pinterf,play,temp,wind1,x_in,K_in,x_surf,cdrag,x_predict)

INTEGER, INTENT(IN) :: ngrid ! number of horizontal index (flat grid)
INTEGER, INTENT(IN) :: nlay ! number of vertical index  
REAL, INTENT(IN) :: rg,rd,dtime ! gravity, R dry air and timestep
REAL, DIMENSION(ngrid,nlay),   INTENT(IN)  :: play      ! pressure middle of layers (Pa)
REAL, DIMENSION(ngrid,nlay),   INTENT(IN)  :: temp      ! temperature (K)
REAL, DIMENSION(ngrid),        INTENT(IN)  :: wind1     ! wind speed first level (m/s)
REAL, DIMENSION(ngrid,nlay+1), INTENT(IN)  :: pinterf ! pressure at interfaces(Pa)
REAL, DIMENSION(ngrid,nlay),   INTENT(IN)  :: x_in    ! variable at the beginning of timestep
REAL, DIMENSION(ngrid,nlay+1), INTENT(IN)  :: K_in    ! eddy diffusivity coef at the beginning of time step
REAL, DIMENSION(ngrid),        INTENT(IN)  :: x_surf  ! surface variable
REAL, DIMENSION(ngrid),        INTENT(IN)  :: cdrag  ! drag coefficient
REAL, DIMENSION(ngrid,nlay),   INTENT(OUT) :: x_predict  ! variable at the end of time step after explicit prediction


integer i,k
real ml,F1,rho
real, dimension(ngrid) :: play1,temp1
real, dimension(ngrid,nlay+1) :: K_big

! computation of K_big

play1(:)=play(:,1)
temp1(:)=temp(:,1)

! "big K" calculation
do  k=2,nlay-1
   do i=1,ngrid
      rho=pinterf(i,k)/rd/(0.5*(temp(i,k-1)+temp(i,k)))
      K_big(i,k)=rg*K_in(i,k)/(play(i,k)-play(i,k+1))*(rho**2)
   enddo
enddo
! speficic treatment for k=nlay
do i=1,ngrid
   rho=pinterf(i,nlay)/rd/temp(i,nlay)
   K_big(i,nlay)=rg*K_in(i,nlay)/(2*(play(i,nlay)-pinterf(i,nlay+1)))*(rho**2)
enddo



! x_predict calculation for 2<=k<=nlay-1
do  k=2,nlay-1
   do i=1,ngrid
      ml=(pinterf(i,k)-pinterf(i,k+1))/rg
      x_predict(i,k)=x_in(i,k)-dtime/ml*(-K_big(i,k+1)*x_in(i,k+1) &
                  + (K_big(i,k)+K_big(i,k+1))*x_in(i,k) &
                  - K_big(i,k)*x_in(i,k-1))
   enddo
enddo

! Specific treatment for k=1 
do i=1,ngrid
   ml=(pinterf(i,1)-pinterf(i,2))/rg
   F1=-play1(i)/rd/temp1(i)*wind1(i)*cdrag(i)*(x_in(i,1)-x_surf(i)) ! attention convention sens du flux
   x_predict(i,1)=x_in(i,1)-dtime/ml*(-K_big(i,2)*(x_in(i,2) - x_in(i,1))-F1)
enddo

! Specific treatment for k=nlay
! flux at the top of the atmosphere=0
do i=1,ngrid
   ml=0.5*(pinterf(i,nlay)-pinterf(i,nlay+1))/rg
   x_predict(i,nlay)=x_in(i,nlay)+dtime/ml*(K_big(i,nlay)*(x_in(i,nlay)-x_in(i,nlay-1)))
enddo


end subroutine atke_explicit_prediction



end module lmdz_call_atke