module math_mod
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose: The module contains all the mathematical SUBROUTINE used in the PEM
!!!
!!! Author: Adapted from Schorgofer MSIM (N.S, Icarus 2010), impletented here by LL
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

implicit none

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

SUBROUTINE deriv1(z,nz,y,y0,ybot,dzY)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose: Compute the first derivative of a function y(z) on an irregular grid
!!!
!!! Author: From N.S (N.S, Icarus 2010), impletented here by LL
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! first derivative of a function y(z) on irregular grid
! upper boundary conditions: y(0) = y0
! lower boundary condition.: yp = ybottom

implicit none

! Inputs
!-------
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
! Outputs
!--------
real, dimension(nz), intent(out) :: dzY ! derivative of y w.r.t depth
! Local variables
!----------------
integer :: j
real    :: hm, hp, c1, c2, c3

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)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose: Compute the second derivative of a function y(z) on an irregular grid
!!!
!!! Author: N.S (raw copy/paste from MSIM)
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! second derivative y_zz on irregular grid
! boundary conditions: y(0) = y0, y(nz + 1) = yNp1

implicit none

! Inputs
!-------
integer,             intent(in) :: nz
real, dimension(nz), intent(in) :: z, y
real,                intent(in) :: y0, yNp1
! Outputs
!--------
real, dimension(nz), intent(out) :: yp2
! Local variables
!----------------
integer :: j
real    :: hm, hp, c1, c2, c3

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)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose: First derivative of function y(z) at z(j)  one-sided derivative on irregular grid
!!!
!!! Author: N.S (raw copy/paste from MSIM)
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

implicit none

! Inputs
!-------
integer,             intent(in) :: nz, j
real, dimension(nz), intent(in) :: z, y
! Outputs
!--------
real, intent(out) :: dy_zj
! Local viariables
!-----------------
real :: h1, h2, c1, c2, c3

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)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose:  Column integrates y on irregular grid
!!!
!!! Author: N.S (raw copy/paste from MSIM)
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

implicit none

! Inputs
!-------
integer,             intent(in) :: nz, i1, i2
real, dimension(nz), intent(in) :: y, z
! Outputs
!--------
real, intent(out) :: integral
! Local viariables
!-----------------
integer             :: i
real, dimension(nz) :: dz

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)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose: Compute the root zr, between two values y1 and y2 at depth z1,z2
!!!
!!! Author: LL
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

implicit none

! Inputs
!-------
real, intent(in) :: y1, y2 ! difference between surface water density and at depth  [kg/m^3]
real, intent(in) :: z1, z2 ! depth [m]
! Outputs
!--------
real, intent(out) :: zr ! depth at which we have zero

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

END SUBROUTINE findroot

!=======================================================================

SUBROUTINE solve_tridiag(a,b,c,d,n,x,error)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose: 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
!!!
!!! Author: JBC
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

implicit none

! Inputs
!-------
integer,                intent(in) :: n
real, dimension(n),     intent(in) :: b, d
real, dimension(n - 1), intent(in) :: a, c
! Outputs
!--------
real, dimension(n), intent(out) :: x
integer,            intent(out) :: error
! Local viariables
!-----------------
integer            :: i
real               :: m
real, dimension(n) :: cp, dp

! 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)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!
!!! Purpose: Solve 1D steady-state heat equation with space-dependent diffusivity
!!!          Left boudary condition : prescribed temperature T_left
!!!          Right boudary condition: prescribed thermal flux q_right
!!!
!!!          z     : grid points
!!!          mz    : mid-grid points
!!!          kappa : thermal diffusivity at grid points
!!!          mkappa: thermal diffusivity at mid-grid points
!!!
!!! Author: JBC
!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

use abort_pem_mod, only: abort_pem

implicit none

! Inputs
!-------
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
! Outputs
!--------
real, dimension(n), intent(out) :: T
! Local viariables
!-----------------
integer                :: i, error
real, dimension(n)     :: b, d
real, dimension(n - 1) :: a, c

! 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 abort_pem("solve_steady_heat","Unstable solving!",1)

END SUBROUTINE solve_steady_heat

end module math_mod