MODULE soil_pem_compute_mod

implicit none

!=======================================================================
contains
!=======================================================================

SUBROUTINE soil_pem_compute(ngrid,nsoil,firstcall,therm_i,timestep,tsurf,tsoil)

use comsoil_h_PEM, only: layer_PEM, mlayer_PEM, mthermdiff_PEM, thermdiff_PEM, coefq_PEM, coefd_PEM, mu_PEM, alph_PEM, beta_PEM, fluxgeo
use comsoil_h,     only: volcapa

implicit none

!-----------------------------------------------------------------------
!  Author: LL
!  Purpose: Compute soil temperature using an implict 1st order scheme
!  
!  Note: depths of layers and mid-layers, soil thermal inertia and 
!        heat capacity are commons in comsoil_PEM.h
!-----------------------------------------------------------------------

#include "dimensions.h"

!-----------------------------------------------------------------------
!  arguments
!  ---------
!  inputs:
integer,                      intent(in) :: ngrid     ! number of (horizontal) grid-points 
integer,                      intent(in) :: nsoil     ! number of soil layers  
logical,                      intent(in) :: firstcall ! identifier for initialization call 
real, dimension(ngrid,nsoil), intent(in) :: therm_i   ! thermal inertia [SI]
real,                         intent(in) :: timestep  ! time step [s]
real, dimension(ngrid),       intent(in) :: tsurf     ! surface temperature [K]
 
! outputs:
real, dimension(ngrid,nsoil), intent(inout) :: tsoil ! soil (mid-layer) temperature [K]
! local variables:
integer :: ig, ik    

! 0. Initialisations and preprocessing step
 if (firstcall) then
! 0.1 Build mthermdiff_PEM(:), the mid-layer thermal diffusivities
    do ig = 1,ngrid
        do ik = 0,nsoil - 1
            mthermdiff_PEM(ig,ik) = therm_i(ig,ik + 1)*therm_i(ig,ik + 1)/volcapa    
        enddo
    enddo

! 0.2 Build thermdiff(:), the "interlayer" thermal diffusivities
    do ig = 1,ngrid
        do ik = 1,nsoil - 1
            thermdiff_PEM(ig,ik) = ((layer_PEM(ik) - mlayer_PEM(ik - 1))*mthermdiff_PEM(ig,ik) &
                                   + (mlayer_PEM(ik) - layer_PEM(ik))*mthermdiff_PEM(ig,ik - 1))/(mlayer_PEM(ik) - mlayer_PEM(ik - 1))
        enddo
    enddo

! 0.3 Build coefficients mu_PEM, q_{k+1/2}, d_k, alph_PEMa_k and capcal
    ! mu_PEM
    mu_PEM = mlayer_PEM(0)/(mlayer_PEM(1) - mlayer_PEM(0))

    ! q_{1/2}
    coefq_PEM(0) = volcapa*layer_PEM(1)/timestep
    ! q_{k+1/2}
    do ik = 1,nsoil - 1
        coefq_PEM(ik) = volcapa*(layer_PEM(ik + 1) - layer_PEM(ik))/timestep
    enddo

    do ig = 1,ngrid
        ! d_k
        do ik = 1,nsoil - 1
            coefd_PEM(ig,ik) = thermdiff_PEM(ig,ik)/(mlayer_PEM(ik)-mlayer_PEM(ik - 1))
        enddo

        ! alph_PEM_{N-1}
        alph_PEM(ig,nsoil - 1) = coefd_PEM(ig,nsoil-1)/(coefq_PEM(nsoil - 1) + coefd_PEM(ig,nsoil - 1))
        ! alph_PEM_k
        do ik = nsoil - 2,1,-1
            alph_PEM(ig,ik) = coefd_PEM(ig,ik)/(coefq_PEM(ik) + coefd_PEM(ig,ik + 1)*(1. - alph_PEM(ig,ik + 1)) + coefd_PEM(ig,ik))
        enddo
      enddo ! of do ig=1,ngrid
endif ! of if (firstcall)

IF (.not. firstcall) THEN
! 2. Compute soil temperatures
! First layer:
    do ig = 1,ngrid
        tsoil(ig,1) = (tsurf(ig) + mu_PEM*beta_PEM(ig,1)*thermdiff_PEM(ig,1)/mthermdiff_PEM(ig,0))/ &
                      (1. + mu_PEM*(1. - alph_PEM(ig,1))*thermdiff_PEM(ig,1)/mthermdiff_PEM(ig,0))

! Other layers:
        do ik = 1,nsoil - 1
                tsoil(ig,ik + 1) = alph_PEM(ig,ik)*tsoil(ig,ik) + beta_PEM(ig,ik) 
        enddo
    enddo
endif

!  2. Compute beta_PEM coefficients (preprocessing for next time step)
! Bottom layer, beta_PEM_{N-1}
do ig = 1,ngrid
    beta_PEM(ig,nsoil - 1) = coefq_PEM(nsoil - 1)*tsoil(ig,nsoil)/(coefq_PEM(nsoil - 1) + coefd_PEM(ig,nsoil - 1)) &
                             + fluxgeo/(coefq_PEM(nsoil - 1) + coefd_PEM(ig,nsoil - 1))
enddo
! Other layers
do ik = nsoil-2,1,-1
    do ig = 1,ngrid
        beta_PEM(ig,ik) = (coefq_PEM(ik)*tsoil(ig,ik + 1) + coefd_PEM(ig,ik + 1)*beta_PEM(ig,ik + 1))/ &
                          (coefq_PEM(ik) + coefd_PEM(ig,ik + 1)*(1. - alph_PEM(ig,ik + 1)) + coefd_PEM(ig,ik))
    enddo
enddo

END SUBROUTINE soil_pem_compute

END MODULE soil_pem_compute_mod