MODULE soil_temp
!-----------------------------------------------------------------------
! NAME
!     soil_temp
!
! DESCRIPTION
!     Routines to compute soil temperature evolution and initialization.
!
! AUTHORS & DATE
!     L. Lange, 2023
!     JB Clement, 2023-2026
!
! NOTES
!     See the header of the 'soil' module for the vertical soil discretization.
!-----------------------------------------------------------------------

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

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

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

!=======================================================================
SUBROUTINE compute_tsoil(ngrid,nsoil,islope,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_deep, volcapa

! 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)    :: islope    ! Index of the sub-slope being computed
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 conductivities
    do ig = 1,ngrid
        do ik = 0,nsoil - 1
            mthermdiff(ig,ik,islope) = therm_i(ig,ik + 1)*therm_i(ig,ik + 1)/volcapa
        end do
    end do

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

! 0.3 Build coefficients q_{k+1/2}, d_k, alph
    ! 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,islope) = thermdiff(ig,ik,islope)/(mlayer(ik)-mlayer(ik - 1))
        end do

        ! alph_PEM_{N-1}
        alph(ig,nsoil - 1,islope) = coefd(ig,nsoil - 1,islope)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1,islope))
        ! alph_PEM_k
        do ik = nsoil - 2,1,-1
            alph(ig,ik,islope) = coefd(ig,ik,islope)/(coefq(ik) + coefd(ig,ik + 1,islope)*(1._dp - alph(ig,ik + 1,islope)) + coefd(ig,ik,islope))
        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,islope)*thermdiff(ig,1,islope)/mthermdiff(ig,0,islope))/ &
                      (1._dp + mu*(1._dp - alph(ig,1,islope))*thermdiff(ig,1,islope)/mthermdiff(ig,0,islope))

! Other layers:
        do ik = 1,nsoil - 1
            tsoil(ig,ik + 1) = alph(ig,ik,islope)*tsoil(ig,ik) + beta(ig,ik,islope)
        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,islope) = coefq(nsoil - 1)*tsoil(ig,nsoil)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1,islope)) &
                                + flux_geo_deep/(coefq(nsoil - 1) + coefd(ig,nsoil - 1,islope))
end do
! Other layers
do ik = nsoil - 2,1,-1
    do ig = 1,ngrid
        beta(ig,ik,islope) = (coefq(ik)*tsoil(ig,ik + 1) + coefd(ig,ik + 1,islope)*beta(ig,ik + 1,islope))/ &
                             (coefq(ik) + coefd(ig,ik + 1,islope)*(1._dp - alph(ig,ik + 1,islope)) + coefd(ig,ik,islope))
    end do
end do

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

!=======================================================================
SUBROUTINE ini_tsoil_pem(ngrid,nsoil,islope,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. Lange, 2023
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! 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)    :: islope  ! Index of the sub-slope being computed
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 conductivities
do ig = 1,ngrid
    do ik = 0,nsoil - 1
        mthermdiff(ig,ik,islope) = therm_i(ig,ik + 1)*therm_i(ig,ik + 1)/volcapa
    end do
end do

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

! 0.3 Build coefficients q_{k+1/2}, d_k, alph
! q_{1/2}
coefq(:) = 0._dp
! q_{k+1/2}

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

    ! alph_PEM_{N-1}
    alph(ig,nsoil - 1,islope) = coefd(ig,nsoil - 1,islope)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1,islope))
    ! alph_PEM_k
    do ik = nsoil - 2,1,-1
        alph(ig,ik,islope) = coefd(ig,ik,islope)/(coefq(ik) + coefd(ig,ik + 1,islope)*(1._dp - alph(ig,ik + 1,islope)) + coefd(ig,ik,islope))
    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,islope) = coefq(nsoil - 1)*tsoil(ig,nsoil)/(coefq(nsoil - 1) + coefd(ig,nsoil - 1,islope)) &
                                    + flux_geo_deep/(coefq(nsoil - 1) + coefd(ig,nsoil - 1,islope))
end do
! Other layers
do ik = nsoil - 2,1,-1
    do ig = 1,ngrid
        beta(ig,ik,islope) = (coefq(ik)*tsoil(ig,ik + 1) + coefd(ig,ik + 1,islope)*beta(ig,ik + 1,islope))/ &
                             (coefq(ik) + coefd(ig,ik + 1,islope)*(1._dp - alph(ig,ik + 1,islope)) + coefd(ig,ik,islope))
    end do
end do

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

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

!=======================================================================
SUBROUTINE shift_tsoil2surf(ngrid,nsoil,nslope,zshift_surf,tsurf,tsoil,zlag)
!-----------------------------------------------------------------------
! NAME
!     shift_tsoil2surf
!
! DESCRIPTION
!     Shift soil temperature profile to follow surface evolution due to ice
!     condensation/sublimation.
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!     Only the planets building a dust lag when their surface ice
!     sublimates uses 'zlag'. When it is absent the routine behaves
!     as if the lag were always too thin to be resolved (ilag = 0).
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use soil,    only: layer, mlayer, flux_geo_deep, thermdiff, get_isoil
use maths,   only: solve_steady_heat
use display, only: print_msg, LVL_NFO

! 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)           :: tsurf       ! Surface temperature [K]
real(dp), dimension(:,:,:), intent(inout)        :: tsoil       ! Soil (mid-layer) temperature [K]
real(dp), dimension(:,:),   intent(in), optional :: zlag        ! Newly built lag thickness [m]

! LOCAL VARIABLES
! ---------------
integer(di)                             :: ig, isoil, islope, ishift, ilag
real(dp)                                :: z, zshift_surfloc, tsoil_top, tsoil_lag_bot
real(dp), dimension(:,:,:), allocatable :: tsoil_old
logical(k4)                             :: has_lag

! CODE
! ----
call print_msg("> Shifting soil temperature profile to match surface evolution",LVL_NFO)
allocate(tsoil_old(ngrid,nsoil,nslope))
tsoil_old = tsoil
has_lag = present(zlag)

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
                ! Nothing of the old column is left: solving the steady heat equation (isothermal at 'tsurf' when 'flux_geo_deep = 0')
                call solve_steady_heat(nsoil,mlayer,layer,thermdiff(ig,:,islope),tsurf(ig,islope),flux_geo_deep,tsoil(ig,:,islope))
            else
                ! Number of nodes buried in the newly deposited material
                ishift = get_isoil(zshift_surfloc)
                ! 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
            ! Deepest node which still maps inside the old column
            ishift = get_isoil(mlayer(nsoil - 1) + zshift_surfloc)
            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) .or. ishift < 1) then ! Surface change is much larger than the discretization
                call solve_steady_heat(nsoil,mlayer,layer,thermdiff(ig,:,islope),tsurf(ig,islope),flux_geo_deep,tsoil(ig,:,islope))
            else
                if (.not. has_lag) then ! This planet does not build a lag
                    ilag = 0
                else if (zlag(ig,islope) < mlayer(0)) then ! The lag is too thin to be taken into account
                    ilag = 0
                else
                    ! Number of nodes buried in the newly built lag
                    ilag = min(get_isoil(zlag(ig,islope)),ishift) ! Capped by 'ishift' so that the lag and the steady-state reconstruction for the "new deepest layers" cannot overlap
                    ! Position of the lag bottom in the old discretization of the depth
                    z = zlag(ig,islope) - zshift_surfloc
                    ! Temperature of the ice left just below the newly built lag
                    tsoil_lag_bot = itp_tsoil(tsoil_old(ig,:,islope),tsurf(ig,islope),z)
                    ! The "new lag" temperature is the steady conduction profile between 'tsurf' and the ice temperature at its bottom
                    do isoil = 1,ilag
                        tsoil(ig,isoil,islope) = tsurf(ig,islope) + (tsoil_lag_bot - tsurf(ig,islope))*min(mlayer(isoil - 1)/zlag(ig,islope),1._dp)
                    end do
                end if

                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
                tsoil_top = tsoil(ig,ishift,islope)
                call solve_steady_heat(nsoil - ishift + 1,mlayer(ishift - 1:),layer(ishift:),thermdiff(ig,ishift:,islope),tsoil_top,flux_geo_deep,tsoil(ig,ishift:,islope))
            end if
        end if
    end do
end do

deallocate(tsoil_old)

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, get_isoil

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

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

! RESULT
! ------
real(dp) :: tsoil_z ! Soil temperature at the depth z [K]

! LOCAL VARIABLES
! ---------------
real(dp)    :: tsoil_1, mlayer_1, a
integer(di) :: iz, imax

! CODE
! ----
! Handle endpoints explicitly
imax = ubound(mlayer,1)
if (z <= 0._dp) then
    tsoil_z = tsurf
    return
end if
if (z >= mlayer(imax)) then
    tsoil_z = tsoil(imax + 1)
    return
end if

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

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

!=======================================================================
SUBROUTINE evolve_soil_temp(tsoil_avg,tsoil_avg_ref,tsurf_avg,tsoil_ts,tsoil_ts_old)
!-----------------------------------------------------------------------
! NAME
!     evolve_soil_temp
!
! DESCRIPTION
!     Update soil temperature profile and its time series.
!
! AUTHORS & DATE
!     JB Clement, 02/2026
!
! NOTES
!     'tsoil_avg_ref' must be the average profile as it was before
!     'shift_tsoil2surf' was applied. Using the shifted profile as the
!     reference instead would add the discontinuity created by the shift
!     into the deviation.
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use geometry,   only: ngrid, nsoil, nslope, nday
use evolution,  only: dt
use soil,       only: TI
use stoppage,   only: stop_clean
use display,    only: print_msg, LVL_NFO

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

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

! LOCAL VARIABLES
! ---------------
integer(di) :: i, isoil, islope, iday

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

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

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

    ! Safety check before the values are used below
    do i = 1,ngrid
        do isoil = 1,nsoil
            if (isnan(tsoil_avg(i,isoil,islope))) call stop_clean(__FILE__,__LINE__,"NaN detected in tsoil_avg",1)
        end do
    end do

    do iday = 1,nday
        do i = 1,ngrid
            do isoil = 1,nsoil
                ! Keep the soil temperature timeseries consistent with the updated daily mean
                ! Tsoil_new = Tsoil_avg_new + Tsoil_dev = Tsoil_avg_new + Tsoil_old - Tsoil_avg_ref
                tsoil_ts(i,isoil,islope,iday) = tsoil_ts(i,isoil,islope,iday) + tsoil_avg(i,isoil,islope) - tsoil_avg_ref(i,isoil,islope)
            end do
        end do
    end do
end do

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

END MODULE soil_temp