module math_toolkit
!-----------------------------------------------------------------------
! NAME
!     math_toolkit
!
! DESCRIPTION
!     The module contains all the mathematical subroutines used in the PEM.
!
! AUTHORS & DATE
!     L. Lange
!     JB Clement, 2023-2025
!
! NOTES
!     Adapted from Schorghofer MSIM (N.S, Icarus 2010)
!-----------------------------------------------------------------------

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

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

!=======================================================================
SUBROUTINE deriv1(z,nz,y,y0,ybot,dzY)
!-----------------------------------------------------------------------
! NAME
!     deriv1
!
! DESCRIPTION
!     Compute the first derivative of a function y(z) on an irregular grid.
!
! AUTHORS & DATE
!     N.S (Icarus 2010)
!     L. Lange
!     JB Clement, 2023-2025
!
! NOTES
!     Upper boundary conditions: y(0) = y0
!     Lower boundary condition: yp = ybottom
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
integer,             intent(in)  :: nz       ! number of layer
real, dimension(nz), intent(in)  :: z        ! depth layer
real, dimension(nz), intent(in)  :: y        ! function which needs to be derived
real,                intent(in)  :: y0, ybot ! boundary conditions
real, dimension(nz), intent(out) :: dzY      ! derivative of y w.r.t depth

! LOCAL VARIABLES
! ---------------
integer :: j
real    :: hm, hp, c1, c2, c3

! CODE
! ----
hp = z(2) - z(1)
c1 = z(1)/(hp*z(2))
c2 = 1/z(1) - 1/(z(2) - z(1))
c3 = -hp/(z(1)*z(2))
dzY(1) = c1*y(2) + c2*y(1) + c3*y0
do j = 2,nz - 1
    hp = z(j + 1) - z(j)
    hm = z(j) - z(j - 1)
    c1 = +hm/(hp*(z(j + 1) - z(j - 1)))
    c2 = 1/hm - 1/hp
    c3 = -hp/(hm*(z(j + 1) - z(j - 1)))
    dzY(j) = c1*y(j + 1) + c2*y(j) + c3*y(j - 1)
enddo
dzY(nz) = (ybot - y(nz - 1))/(2.*(z(nz) - z(nz - 1)))

END SUBROUTINE deriv1
!=======================================================================

!=======================================================================
SUBROUTINE deriv2_simple(z,nz,y,y0,yNp1,yp2)
!-----------------------------------------------------------------------
! NAME
!     deriv2_simple
!
! DESCRIPTION
!     Compute the second derivative of a function y(z) on an irregular grid.
!
! AUTHORS & DATE
!     N.S (Icarus 2010)
!     JB Clement, 2023-2025
!
! NOTES
!     Boundary conditions: y(0) = y0, y(nz + 1) = yNp1
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
integer,             intent(in)  :: nz
real, dimension(nz), intent(in)  :: z, y
real,                intent(in)  :: y0, yNp1
real, dimension(nz), intent(out) :: yp2

! LOCAL VARIABLES
! ---------------
integer :: j
real    :: hm, hp, c1, c2, c3

! CODE
! ----
c1 = +2./((z(2) - z(1))*z(2))
c2 = -2./((z(2) - z(1))*z(1))
c3 = +2./(z(1)*z(2))
yp2(1) = c1*y(2) + c2*y(1) + c3*y0
do j = 2,nz - 1
    hp = z(j + 1) - z(j)
    hm = z(j) - z(j - 1)
    c1 = +2./(hp*(z(j + 1) - z(j - 1)))
    c2 = -2./(hp*hm)
    c3 = +2./(hm*(z(j + 1) - z(j-1)))
    yp2(j) = c1*y(j + 1) + c2*y(j) + c3*y(j - 1)
enddo
yp2(nz) = (yNp1 - 2*y(nz) + y(nz - 1))/(z(nz) - z(nz - 1))**2

END SUBROUTINE deriv2_simple
!=======================================================================

!=======================================================================
SUBROUTINE  deriv1_onesided(j,z,nz,y,dy_zj)
!-----------------------------------------------------------------------
! NAME
!     deriv1_onesided
!
! DESCRIPTION
!     First derivative of function y(z) at z(j) one-sided derivative
!     on irregular grid.
!
! AUTHORS & DATE
!     N.S (Icarus 2010)
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
integer,             intent(in)  :: nz, j
real, dimension(nz), intent(in)  :: z, y
real,                intent(out) :: dy_zj

! LOCAL VARIABLES
! ---------------
real :: h1, h2, c1, c2, c3

! CODE
! ----
if (j < 1 .or. j > nz - 2) then
    dy_zj = -1.
else
    h1 = z(j + 1) - z(j)
    h2 = z(j + 2)- z(j + 1)
    c1 = -(2*h1 + h2)/(h1*(h1 + h2))
    c2 = (h1 + h2)/(h1*h2)
    c3 = -h1/(h2*(h1 + h2))
    dy_zj = c1*y(j) + c2*y(j + 1) + c3*y(j + 2)
endif

END SUBROUTINE deriv1_onesided
!=======================================================================

!=======================================================================
PURE SUBROUTINE colint(y,z,nz,i1,i2,integral)
!-----------------------------------------------------------------------
! NAME
!     colint
!
! DESCRIPTION
!     Column integrates y on irregular grid.
!
! AUTHORS & DATE
!     N.S (Icarus 2010)
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
integer,             intent(in)  :: nz, i1, i2
real, dimension(nz), intent(in)  :: y, z
real,                intent(out) :: integral

! LOCAL VARIABLES
! ---------------
integer             :: i
real, dimension(nz) :: dz

! CODE
! ----
dz(1) = (z(2) - 0.)/2
do i = 2,nz - 1
    dz(i) = (z(i + 1) - z(i - 1))/2.
enddo
dz(nz) = z(nz) - z(nz - 1)
integral = sum(y(i1:i2)*dz(i1:i2))

END SUBROUTINE colint
!=======================================================================

!=======================================================================
SUBROUTINE findroot(y1,y2,z1,z2,zr)
!-----------------------------------------------------------------------
! NAME
!     findroot
!
! DESCRIPTION
!     Compute the root zr, between two values y1 and y2 at depth z1,z2.
!
! AUTHORS & DATE
!     L. Lange
!     JB Clement, 2023-2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
real, intent(in)  :: y1, y2
real, intent(in)  :: z1, z2
real, intent(out) :: zr

! CODE
! ----
zr = (y1*z2 - y2*z1)/(y1 - y2)

END SUBROUTINE findroot
!=======================================================================

!=======================================================================
SUBROUTINE solve_tridiag(a,b,c,d,n,x,error)
!-----------------------------------------------------------------------
! NAME
!     solve_tridiag
!
! DESCRIPTION
!     Solve a tridiagonal system Ax = d using the Thomas' algorithm
!     a: sub-diagonal
!     b: main diagonal
!     c: super-diagonal
!     d: right-hand side
!     x: solution
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!
!-----------------------------------------------------------------------

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

! ARGUMENTS
! ---------
integer,                intent(in)  :: n
real, dimension(n),     intent(in)  :: b, d
real, dimension(n - 1), intent(in)  :: a, c
real, dimension(n),     intent(out) :: x
integer,                intent(out) :: error

! LOCAL VARIABLES
! ---------------
integer            :: i
real               :: m
real, dimension(n) :: cp, dp

! CODE
! ----
! Check stability: diagonally dominant condition
error = 0
if (abs(b(1)) < abs(c(1))) then
    error = 1
    return
endif
do i = 2,n - 1
    if (abs(b(i)) < abs(a(i - 1)) + abs(c(i))) then
        error = 1
        return
    endif
enddo
if (abs(b(n)) < abs(a(n - 1))) then
    error = 1
    return
endif

! Initialization
cp(1) = c(1)/b(1)
dp(1) = d(1)/b(1)

! Forward sweep
do i = 2,n - 1
    m = b(i) - a(i - 1)*cp(i - 1)
    cp(i) = c(i)/m
    dp(i) = (d(i) - a(i - 1)*dp(i - 1))/m
enddo
m = b(n) - a(n - 1)*cp(n - 1)
dp(n) = (d(n) - a(n - 1)*dp(n - 1))/m

! Backward substitution
x(n) = dp(n)
do i = n - 1,1,-1
    x(i) = dp(i) - cp(i)*x(i + 1)
enddo

END SUBROUTINE solve_tridiag
!=======================================================================

!=======================================================================
SUBROUTINE solve_steady_heat(n,z,mz,kappa,mkappa,T_left,q_right,T)
!-----------------------------------------------------------------------
! NAME
!     solve_steady_heat
!
! DESCRIPTION
!     Solve 1D steady-state heat equation with space-dependent thermal diffusivity.
!     Uses Thomas algorithm to solve tridiagonal system.
!     Left boundary: prescribed temperature (Dirichlet).
!     Right boundary: prescribed thermal flux (Neumann).
!
! AUTHORS & DATE
!     JB Clement, 2025
!
! NOTES
!     Grid points at z, mid-grid points at mz.
!     Thermal diffusivity specified at both grids (kappa, mkappa).
!-----------------------------------------------------------------------

! DEPENDENCIES
! ------------
use stop_pem, only: stop_clean

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

! ARGUMENTS
! ---------
integer,                intent(in)  :: n
real, dimension(n),     intent(in)  :: z, kappa
real, dimension(n - 1), intent(in)  :: mz, mkappa
real,                   intent(in)  :: T_left, q_right
real, dimension(n),     intent(out) :: T

! LOCAL VARIABLES
! ---------------
integer                :: i, error
real, dimension(n)     :: b, d
real, dimension(n - 1) :: a, c

! CODE
! ----
! Initialization
a = 0.; b = 0.; c = 0.; d = 0.

! Left boundary condition (Dirichlet: prescribed temperature)
b(1) = 1.
d(1) = T_left

! Internal points
do i = 2,n - 1
    a(i - 1) = -mkappa(i - 1)/((mz(i) - mz(i - 1))*(z(i) - z(i - 1)))
    c(i) = -mkappa(i)/((mz(i) - mz(i - 1))*(z(i + 1) - z(i)))
    b(i) = -(a(i - 1) + c(i))
enddo

! Right boundary condition (Neumann: prescribed temperature)
a(n - 1) = kappa(n - 1)/(z(n) - z(n - 1))
b(n) = -kappa(n)/(z(n) - z(n - 1))
d(n) = q_right

! Solve the tridiagonal linear system with the Thomas' algorithm
call solve_tridiag(a,b,c,d,n,T,error)
if (error /= 0) call stop_clean("solve_steady_heat","Unstable solving!",1)

END SUBROUTINE solve_steady_heat
!=======================================================================

end module math_toolkit