MODULE soil_temp
!-----------------------------------------------------------------------
! NAME
!     soil_temp
!
! DESCRIPTION
!     Routines to compute soil temperature evolution and initialization.
!
! AUTHORS & DATE
!     L. Lange, 2023
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use numerics, only: dp, di, k4

! DECLARATION
! -----------
implicit none

! PARAMETERS
! ----------
real(dp), dimension(:,:,:), allocatable, protected :: tsoil_PCM    ! Soil temperature in the PCM at the beginning [K]
real(dp), dimension(:,:),   allocatable, protected :: flux_geo_PCM ! Geothermal flux in the PCM at the beginning [K]

contains
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

!=======================================================================
SUBROUTINE ini_soil_temp()
!-----------------------------------------------------------------------
! NAME
!     ini_soil_temp
!
! DESCRIPTION
!     Initialize the parameters of module 'soil_temp'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry, only: ngrid, nsoil_PCM, nslope

! DECLARATION
! -----------
implicit none

! CODE
! ----
if (.not. allocated(tsoil_PCM)) allocate(tsoil_PCM(ngrid,nsoil_PCM,nslope))
if (.not. allocated(flux_geo_PCM)) allocate(flux_geo_PCM(ngrid,nslope))

END SUBROUTINE ini_soil_temp
!=======================================================================

!=======================================================================
SUBROUTINE end_soil_temp()
!-----------------------------------------------------------------------
! NAME
!     end_soil_temp
!
! DESCRIPTION
!     Deallocate soil_temp arrays.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DECLARATION
! -----------
implicit none

! CODE
! ----
if (allocated(tsoil_PCM)) deallocate(tsoil_PCM)
if (allocated(flux_geo_PCM)) deallocate(flux_geo_PCM)

END SUBROUTINE end_soil_temp
!=======================================================================

!=======================================================================
SUBROUTINE set_tsoil_PCM(tsoil_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_tsoil_PCM
!
! DESCRIPTION
!     Setter for 'tsoil_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
real(dp), dimension(:,:,:), intent(in) :: tsoil_PCM_in

! CODE
! ----
tsoil_PCM(:,:,:) = tsoil_PCM_in(:,:,:)

END SUBROUTINE set_tsoil_PCM
!=======================================================================

!=======================================================================
SUBROUTINE set_flux_geo_PCM(flux_geo_PCM_in)
!-----------------------------------------------------------------------
! NAME
!     set_flux_geo_PCM
!
! DESCRIPTION
!     Setter for 'flux_geo_PCM'.
!
! AUTHORS & DATE
!     JB Clement, 12/2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
real(dp), dimension(:,:), intent(in) :: flux_geo_PCM_in

! CODE
! ----
flux_geo_PCM(:,:) = flux_geo_PCM_in(:,:)

END SUBROUTINE set_flux_geo_PCM
!=======================================================================

!=======================================================================
SUBROUTINE compute_tsoil(ngrid,nsoil,firstcall,therm_i,timestep,tsurf,tsoil)
!-----------------------------------------------------------------------
! NAME
!     compute_tsoil
!
! DESCRIPTION
!     Compute soil temperature using an implicit 1st order scheme.
!
! AUTHORS & DATE
!     L. Lange, 2023
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use soil, only: layer, mlayer, mthermdiff, thermdiff, coefq, coefd, mu, alph, beta, flux_geo, volcapa

! DECLARATION
! -----------
implicit none

#include "dimensions.h"

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

! LOCAL VARIABLES
! ---------------
integer(di) :: ig, ik

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

! 0.2 Build thermdiff(:), the "interlayer" thermal diffusivities
    do ig = 1,ngrid
        do ik = 1,nsoil - 1
            thermdiff(ig,ik) = ((layer(ik) - mlayer(ik - 1))*mthermdiff(ig,ik) &
                                   + (mlayer(ik) - layer(ik))*mthermdiff(ig,ik - 1))/(mlayer(ik) - mlayer(ik - 1))
        end do
    end do

! 0.3 Build coefficients mu, q_{k+1/2}, d_k, alph
    ! mu
    mu = mlayer(0)/(mlayer(1) - mlayer(0))

    ! q_{1/2}
    coefq(0) = volcapa*layer(1)/timestep
    ! q_{k+1/2}
    do ik = 1,nsoil - 1
        coefq(ik) = volcapa*(layer(ik + 1) - layer(ik))/timestep
    end do

    do ig = 1,ngrid
        ! d_k
        do ik = 1,nsoil - 1
            coefd(ig,ik) = thermdiff(ig,ik)/(mlayer(ik)-mlayer(ik - 1))
        end do

        ! alph_PEM_{N-1}
        alph(ig,nsoil - 1) = coefd(ig,nsoil - 1)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1))
        ! alph_PEM_k
        do ik = nsoil - 2,1,-1
            alph(ig,ik) = coefd(ig,ik)/(coefq(ik) + coefd(ig,ik + 1)*(1._dp - alph(ig,ik + 1)) + coefd(ig,ik))
        end do
      end do ! of do ig=1,ngrid
end if ! of if (firstcall)

if (.not. firstcall) THEN
! 2. Compute soil temperatures
! First layer:
    do ig = 1,ngrid
        tsoil(ig,1) = (tsurf(ig) + mu*beta(ig,1)*thermdiff(ig,1)/mthermdiff(ig,0))/ &
                      (1._dp + mu*(1._dp - alph(ig,1))*thermdiff(ig,1)/mthermdiff(ig,0))

! Other layers:
        do ik = 1,nsoil - 1
                tsoil(ig,ik + 1) = alph(ig,ik)*tsoil(ig,ik) + beta(ig,ik)
        end do
    end do
end if

!  2. Compute beta coefficients (preprocessing for next time step)
! Bottom layer, beta_PEM_{N-1}
do ig = 1,ngrid
    beta(ig,nsoil - 1) = coefq(nsoil - 1)*tsoil(ig,nsoil)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1)) &
                             + flux_geo/(coefq(nsoil - 1) + coefd(ig,nsoil - 1))
end do
! Other layers
do ik = nsoil - 2,1,-1
    do ig = 1,ngrid
        beta(ig,ik) = (coefq(ik)*tsoil(ig,ik + 1) + coefd(ig,ik + 1)*beta(ig,ik + 1))/ &
                          (coefq(ik) + coefd(ig,ik + 1)*(1._dp - alph(ig,ik + 1)) + coefd(ig,ik))
    end do
end do

END SUBROUTINE compute_tsoil
!=======================================================================

!=======================================================================
SUBROUTINE ini_tsoil_pem(ngrid,nsoil,therm_i,tsurf,tsoil)
!-----------------------------------------------------------------------
! NAME
!     ini_tsoil_pem
!
! DESCRIPTION
!     Initialize soil with the solution of the stationary heat conduction problem.
!     Boundary conditions: Tsurf averaged from PCM; geothermal flux at bottom.
!
! AUTHORS & DATE
!     L. Lang, 2023
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use soil, only: layer, mlayer, mthermdiff, thermdiff, coefq, coefd, mu, alph, beta, flux_geo, volcapa

! DECLARATION
! -----------
implicit none

#include "dimensions.h"

! ARGUMENTS
! ---------
integer(di),              intent(in)    :: ngrid   ! number of (horizontal) grid-points
integer(di),              intent(in)    :: nsoil   ! number of soil layers
real(dp), dimension(:,:), intent(in)    :: therm_i ! thermal inertia [SI]
real(dp), dimension(:),   intent(in)    :: tsurf   ! surface temperature [K]
real(dp), dimension(:,:), intent(inout) :: tsoil   ! soil (mid-layer) temperature [K]

! LOCAL VARIABLES
! ---------------
integer(di) :: ig, ik, iloop

! CODE
! ----
! 0. Initialisations and preprocessing step
! 0.1 Build mthermdiff(:), the mid-layer thermal diffusivities
do ig = 1,ngrid
    do ik = 0,nsoil - 1
        mthermdiff(ig,ik) = therm_i(ig,ik + 1)*therm_i(ig,ik + 1)/volcapa
    end do
end do

! 0.2 Build thermdiff(:), the "interlayer" thermal diffusivities
do ig = 1,ngrid
    do ik = 1,nsoil - 1
        thermdiff(ig,ik) = ((layer(ik) - mlayer(ik - 1))*mthermdiff(ig,ik) &
                             + (mlayer(ik) - layer(ik))*mthermdiff(ig,ik - 1))/(mlayer(ik) - mlayer(ik - 1))
    end do
end do

! 0.3 Build coefficients mu, q_{k+1/2}, d_k, alph
! mu
mu = mlayer(0)/(mlayer(1) - mlayer(0))

! q_{1/2}
coefq(:) = 0.
! q_{k+1/2}

do ig = 1,ngrid
    ! d_k
    do ik = 1,nsoil - 1
        coefd(ig,ik) = thermdiff(ig,ik)/(mlayer(ik) - mlayer(ik - 1))
    end do

    ! alph_PEM_{N-1}
    alph(ig,nsoil - 1) = coefd(ig,nsoil - 1)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1))
    ! alph_PEM_k
    do ik = nsoil - 2,1,-1
        alph(ig,ik) = coefd(ig,ik)/(coefq(ik) + coefd(ig,ik + 1)*(1._dp - alph(ig,ik + 1)) + coefd(ig,ik))
    end do
end do ! of do ig=1,ngrid

!  1. Compute beta coefficients
! Bottom layer, beta_PEM_{N-1}
do ig = 1,ngrid
        beta(ig,nsoil - 1) = coefq(nsoil - 1)*tsoil(ig,nsoil)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1)) &
                                 + flux_geo/(coefq(nsoil - 1) + coefd(ig,nsoil - 1))
end do
! Other layers
do ik = nsoil - 2,1,-1
    do ig = 1,ngrid
        beta(ig,ik) = (coefq(ik)*tsoil(ig,ik + 1) + coefd(ig,ik + 1)*beta(ig,ik + 1))/ &
                          (coefq(ik) + coefd(ig,ik + 1)*(1._dp - alph(ig,ik + 1)) + coefd(ig,ik))
    end do
end do

!  2. Compute soil temperatures
do iloop = 1,10 !just convergence
    do ig = 1,ngrid
        ! First layer:
        tsoil(ig,1) = (tsurf(ig) + mu*beta(ig,1)*thermdiff(ig,1)/mthermdiff(ig,0))/ &
                      (1._dp + mu*(1._dp - alph(ig,1))*thermdiff(ig,1)/mthermdiff(ig,0))
        ! Other layers:
        do ik = 1,nsoil - 1
            tsoil(ig,ik + 1) = alph(ig,ik)*tsoil(ig,ik) + beta(ig,ik)
        end do
    end do
end do ! iloop

END SUBROUTINE ini_tsoil_pem
!=======================================================================

!=======================================================================
SUBROUTINE shift_tsoil2surf(ngrid,nsoil,nslope,zshift_surf,zlag,tsurf,tsoil)
!-----------------------------------------------------------------------
! NAME
!     shift_tsoil2surf
!
! DESCRIPTION
!     Shift soil temperature profile to follow surface evolution due to ice
!     condensation/sublimation.
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use soil,    only: layer, mlayer, flux_geo, thermdiff, mthermdiff
use maths,   only: solve_steady_heat
use display, only: print_msg

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
integer(di),                intent(in)    :: ngrid       ! number of (horizontal) grid-points
integer(di),                intent(in)    :: nsoil       ! number of soil layers
integer(di),                intent(in)    :: nslope      ! number of sub-slopes
real(dp), dimension(:,:),   intent(in)    :: zshift_surf ! elevation shift for the surface [m]
real(dp), dimension(:,:),   intent(in)    :: zlag        ! newly built lag thickness [m]
real(dp), dimension(:,:),   intent(in)    :: tsurf       ! surface temperature [K]
real(dp), dimension(:,:,:), intent(inout) :: tsoil       ! soil (mid-layer) temperature [K]

! LOCAL VARIABLES
! ---------------
integer(di)                             :: ig, isoil, islope, ishift, ilag
real(dp)                                :: z, zshift_surfloc
real(dp), dimension(ngrid,nsoil,nslope) :: tsoil_old

! CODE
! ----
call print_msg("> Shifting soil temperature profile to match surface evolution")
tsoil_old = tsoil

do ig = 1,ngrid
    do islope = 1,nslope
        zshift_surfloc = zshift_surf(ig,islope)

        if (zshift_surfloc >= 0.) then ! In case of the surface is higher than initially
            if (zshift_surfloc < mlayer(0)) then ! Surface change is too small to be taken into account
                ! Nothing to do; we keep the soil temperature profile
            else if (zshift_surfloc >= mlayer(nsoil - 1)) then ! Surface change is much larger than the discretization
                tsoil(ig,:,islope) = tsurf(ig,islope)
            else
                ishift = 0
                do while (mlayer(ishift) <= zshift_surfloc)
                    ishift = ishift + 1 ! mlayer indices begin at 0 so this the good index for tsoil!
                end do
                ! The "new soil" temperature is set to tsurf
                tsoil(ig,:ishift,islope) = tsurf(ig,islope)
                do isoil = ishift + 1,nsoil
                    ! Position in the old discretization of the depth
                    z = mlayer(isoil - 1) - zshift_surfloc
                    ! Interpolation of the temperature profile from the old discretization
                    tsoil(ig,isoil,islope) = itp_tsoil(tsoil_old(ig,:,islope),tsurf(ig,islope),z)
                end do
            end if
        else ! In case of the surface is lower than initially
            if (abs(zshift_surfloc) < mlayer(0)) then ! Surface change is too small to be taken into account
                ! Nothing to do; we keep the soil temperature profile
            else if (abs(zshift_surfloc) >= mlayer(nsoil - 1)) then ! Surface change is much larger than the discretization
                call solve_steady_heat(nsoil,mlayer,layer,mthermdiff(ig,:),thermdiff(ig,:),tsurf(ig,islope),flux_geo,tsoil(ig,:,islope))
            else
                if (zlag(ig,islope) < mlayer(0)) then ! The lag is too thin to be taken into account
                    ilag = 0
                else
                    ilag = 0
                    do while (mlayer(ilag) <= zlag(ig,islope))
                        ilag = ilag + 1 ! mlayer indices begin at 0 so this the good index for tsoil!
                    end do
                    ! Position of the lag bottom in the old discretization of the depth
                    z = zlag(ig,islope) - zshift_surfloc
                    ! The "new lag" temperature is set to the ice temperature just below
                    tsoil(ig,:ilag,islope) = itp_tsoil(tsoil_old(ig,:,islope),tsurf(ig,islope),z)
                end if

                ishift = nsoil - 1
                z = mlayer(nsoil - 1) + zshift_surfloc
                do while (mlayer(ishift) >= z)
                    ishift = ishift - 1
                end do
                ishift = ishift + 1 ! Adding 1 is needed to match the good index for tsoil!
                do isoil = ilag + 1,ishift
                    ! Position in the old discretization of the depth
                    z = mlayer(isoil - 1) - zshift_surfloc
                    ! Interpolation of the temperature profile from the old discretization
                    tsoil(ig,isoil,islope) = itp_tsoil(tsoil_old(ig,:,islope),tsurf(ig,islope),z)
                end do

                ! The "new deepest layers" temperature is set by solving the steady heat equation
                call solve_steady_heat(nsoil - ishift + 1,mlayer(ishift - 1:),layer(ishift:),mthermdiff(ig,ishift - 1:),thermdiff(ig,ishift:),tsoil(ig,ishift,islope),flux_geo,tsoil(ig,ishift:,islope))
            end if
        end if
    end do
end do

END SUBROUTINE shift_tsoil2surf
!=======================================================================

!=======================================================================
FUNCTION itp_tsoil(tsoil,tsurf,z) RESULT(tsoil_z)
!-----------------------------------------------------------------------
! NAME
!     itp_tsoil
!
! DESCRIPTION
!     Interpolate soil temperature profile.
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use soil, only: mlayer

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
real(dp), dimension(:), intent(in) :: tsoil
real(dp),               intent(in) :: z, tsurf

! LOCAL VARIABLES
! ---------------
real(dp)    :: tsoil_z, tsoil_minus, mlayer_minus, a
integer(di) :: iz

! CODE
! ----
! Find the interval [mlayer(iz - 1),mlayer(iz)[ where the position z belongs
iz = 0
do while (mlayer(iz) <= z)
    iz = iz + 1
end do
if (iz == 0) then
    tsoil_minus = tsurf
    mlayer_minus = 0._dp
else
    tsoil_minus = tsoil(iz)
    mlayer_minus = mlayer(iz - 1)
end if
! Interpolation of the temperature profile from the old discretization
a = (tsoil(iz + 1) - tsoil_minus)/(mlayer(iz) - mlayer_minus)
tsoil_z = a*(z - mlayer_minus) + tsoil_minus

END FUNCTION itp_tsoil
!=======================================================================

!=======================================================================
SUBROUTINE evolve_soil_temp(tsoil_avg, tsurf_avg, tsoil_ts, tsoil_ts_old, h2o_soildensity_avg)
!-----------------------------------------------------------------------
! NAME
!     evolve_soil_temp
!
! DESCRIPTION
!     Update soil temperature profile and calculate water soil density 
!     time series and averages.
!
! AUTHORS & DATE
!     JB Clement, 02/2026
!
! NOTES
!
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,  only: ngrid, nsoil, nslope, nday
use planet,    only: alpha_clap_h2o, beta_clap_h2o
use tracers,   only: mmol, iPCM_qh2o
use physics,   only: mugaz, r
use evolution, only: dt
use soil,      only: TI
use stoppage,  only: stop_clean
use display,   only: print_msg

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
real(dp), dimension(:,:),     intent(in)    :: tsurf_avg
real(dp), dimension(:,:,:),   intent(inout) :: tsoil_avg
real(dp), dimension(:,:,:,:), intent(inout) :: tsoil_ts
real(dp), dimension(:,:,:,:), intent(out)   :: tsoil_ts_old
real(dp), dimension(:,:,:),   intent(out)   :: h2o_soildensity_avg

! LOCAL VARIABLES
! ---------------
real(dp), dimension(ngrid,nsoil)             :: tsoil_avg_old
real(dp), dimension(ngrid,nsoil,nslope,nday) :: h2o_soildensity_ts
integer(di)                                  :: i, isoil, islope, iday

! CODE
! ----
call print_msg("> Update the soil temperature")

! Store current state
tsoil_ts_old(:,:,:,:) = tsoil_ts(:,:,:,:)

do islope = 1,nslope
    ! Store the average profile before the update
    tsoil_avg_old(:,:) = tsoil_avg(:,:,islope)
    
    ! Compute the new soil temperature
    call compute_tsoil(ngrid,nsoil,.true., TI(:,:,islope),dt,tsurf_avg(:,islope),tsoil_avg(:,:,islope))
    call compute_tsoil(ngrid,nsoil,.false.,TI(:,:,islope),dt,tsurf_avg(:,islope),tsoil_avg(:,:,islope))

    do iday = 1,nday
        do i = 1,ngrid
            do isoil = 1,nsoil
                ! Update soil temperature timeseries which is needed to compute the water soil density timeseries
                tsoil_ts(i,isoil,islope,iday) = tsoil_ts(i,isoil,islope,iday)*tsoil_avg(i,isoil,islope)/tsoil_avg_old(i,isoil)
                ! Update water soil density
                h2o_soildensity_ts(i,isoil,islope,iday) = exp(beta_clap_h2o/tsoil_ts(i,isoil,islope,iday) + alpha_clap_h2o)/tsoil_ts(i,isoil,islope,iday)**mmol(iPCM_qh2o)/(mugaz*r)
                ! Safety check
                if (isnan(tsoil_avg(i,isoil,islope))) call stop_clean(__FILE__,__LINE__,"NaN detected in tsoil_avg",1)
            end do
        end do
    end do
end do

! Calculate the time-averaged water soil density
h2o_soildensity_avg = sum(h2o_soildensity_ts,4)/real(nday,dp)

END SUBROUTINE evolve_soil_temp
!=======================================================================

END MODULE soil_temp