! Module  : math_m
! Purpose : various mathematical functions
!
! Copyright 2021 Rich Townsend & The MSG Team
!
! This file is part of MSG. MSG is free software: you can redistribute
! it and/or modify it under the terms of the GNU General Public
! License as published by the Free Software Foundation, version 3.
!
! MSG is distributed in the hope that it will be useful, but WITHOUT
! ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
! or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
! License for more details.
!
! You should have received a copy of the GNU General Public License
! along with this program.  If not, see <http://www.gnu.org/licenses/>.

#:include 'forum.inc'

module math_m

   ! Uses

   use forum_m

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Access specifiers

   private

   public :: integrate_con
   public :: integrate_lin

   ! Procedures

contains

   function integrate_con(x, y, x_a, x_b) result (iy)

      real(RD), intent(in) :: x(:)
      real(RD), intent(in) :: y(:)
      real(RD), intent(in) :: x_a
      real(RD), intent(in) :: x_b
      real(RD)             :: iy

      integer  :: i_a
      integer  :: i_b
      integer  :: n_x
      real(RD) :: u_a
      real(RD) :: u_b
      real(RD) :: iy_a
      real(RD) :: iy_b

      @:CHECK_BOUNDS(SIZE(y), SIZE(x)-1)

      ! Integrate the piecewise-constant function y(x) from x_a to x_b

      call locate(x, x_a, i_a)
      call locate(x, x_b, i_b)

      n_x = SIZE(x)

      i_a = MIN(i_a, n_x-1)
      i_b = MIN(i_b, n_x-1)

      @:ASSERT_DEBUG(i_a >= 1, 'out-of-bounds integration')
      @:ASSERT_DEBUG(i_a < n_x, 'out-of-bounds integration')

      @:ASSERT_DEBUG(i_b >= 1, 'out-of-bounds integration')
      @:ASSERT_DEBUG(i_b < n_x, 'out-of-bounds integration')

      u_a = (x_a - x(i_a))/(x(i_a+1)-x(i_a))
      u_b = (x_b - x(i_b))/(x(i_b+1)-x(i_b))

      if (i_a == i_b) then

         iy = y(i_a)*(u_b - u_a)*(x(i_a+1) - x(i_a))

      else

         iy_a = y(i_a)*(1._RD - u_a)*(x(i_a+1) - x(i_a))

         iy_b = y(i_b)*u_b*(x(i_b+1) - x(i_b))

         iy = (iy_a + iy_b) + SUM(y(i_a+1:i_b-1)*(x(i_a+2:i_b) - x(i_a+1:i_b-1)))

      endif

      ! Finish

      return

   end function integrate_con

   !****

   function integrate_lin(x, y, x_a, x_b) result (iy)

      real(RD), intent(in) :: x(:)
      real(RD), intent(in) :: y(:)
      real(RD), intent(in) :: x_a
      real(RD), intent(in) :: x_b
      real(RD)             :: iy

      integer  :: i_a
      integer  :: i_b
      integer  :: n_x
      real(RD) :: u_a
      real(RD) :: u_b
      real(RD) :: iy_a
      real(RD) :: iy_b

      @:CHECK_BOUNDS(SIZE(y), SIZE(x))

      ! Integrate the piecewise linear function y(x) from x_a to x_b

      call locate(x, x_a, i_a)
      call locate(x, x_b, i_b)

      n_x = SIZE(x)

      i_a = MIN(i_a, n_x-1)
      i_b = MIN(i_b, n_x-1)

      @:ASSERT_DEBUG(i_a >= 1, 'out-of-bounds integration')
      @:ASSERT_DEBUG(i_a < n_x, 'out-of-bounds integration')

      @:ASSERT_DEBUG(i_b >= 1, 'out-of-bounds integration')
      @:ASSERT_DEBUG(i_b < n_x, 'out-of-bounds integration')

      u_a = (x_a - x(i_a))/(x(i_a+1)-x(i_a))
      u_b = (x_b - x(i_b))/(x(i_b+1)-x(i_b))

      if (i_a == i_b) then

         iy = 0.5_RD*( &
              y(i_a)*(2._RD - u_a - u_b) + y(i_a+1)*(u_a + u_b) ) * &
              (u_b - u_a)*(x(i_a+1) - x(i_a))

      else

         iy_a = 0.5_RD*( &
              y(i_a)*(1._RD - u_a) + y(i_a+1)*(1._RD + u_a) ) * &
              (1._RD - u_a)*(x(i_a+1) - x(i_a))

         iy_b = 0.5_RD*( &
              y(i_b)*(2._RD - u_b) + y(i_b+1)*u_b ) * &
              u_b*(x(i_b+1) - x(i_b))

         iy = (iy_a + iy_b) + 0.5_RD*SUM((y(i_a+1:i_b-1) + y(i_a+2:i_b))*(x(i_a+2:i_b) - x(i_a+1:i_b-1)))

      endif

      ! Finish

      return

   end function integrate_lin

end module math_m
