! Module  : ninterp_m
! Purpose : low-level N-dimensional interpolation support routines
!
! Copyright 2021-2023 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 ninterp_m

   ! Uses

   use forum_m

   use stat_m

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Access specifiers

   private

   public :: eval_deriv_kernel
   public :: eval_interp_kernel

   ! Procedures

contains

   subroutine eval_deriv_kernel(rank, mask, vdelta, K_D, stat, order)

      integer, intent(in)            :: rank
      logical, intent(in)            :: mask(:)
      real(RD), intent(in)           :: vdelta(:,:)
      real(RD), intent(out)          :: K_D(:,:)
      integer, intent(out), optional :: stat
      integer, intent(in), optional  :: order

      integer  :: order_
      integer  :: n_match
      integer  :: i
      integer  :: j
      integer  :: bits_l
      integer  :: bits_r
      logical  :: mask_try(SIZE(mask))
      logical  :: dim_mask(4,rank)
      logical  :: dim_mask_try(4,rank)
      integer  :: n_match_try
      real(RD) :: dim_K(4)

      @:CHECK_BOUNDS(SIZE(mask), 4**rank)

      @:CHECK_BOUNDS(SIZE(vdelta, 1), 3)
      @:CHECK_BOUNDS(SIZE(vdelta, 2), rank)

      @:CHECK_BOUNDS(SIZE(K_D, 1), 4**rank)
      @:CHECK_BOUNDS(SIZE(K_D, 2), 4**rank)

      if (PRESENT(order)) then
         order_ = order
      else
         order_ = 3
      end if

      ! Evaluate the kernel matrix K_D for setting up function values
      ! and derivatives. Given a vector F of function values on a
      ! 4**rank-point stencil, MATMUL(K_D, F) gives the function values
      ! and derivatives on the center 2**rank-point stencil

      ! First set up the rank-dimensional mask that defines which
      ! stencil points will be used by K_D. This mask is built from
      ! the inner product of per-dimension masks

      n_match = 0

      select case(order_)

      case(1)

         ! Linear case: set per-dimension masks to include only
         ! interior points

         ! Accumulate the outer product of per-dimension masks

         mask_try(1) = .TRUE.

         do j = 1, rank

            ! Set up the per-dimension mask

            dim_mask_try(1,j) = .FALSE.
            dim_mask_try(2:3,j) = .TRUE.
            dim_mask_try(4,j) = .FALSE.

            ! Accumulate it

            mask_try(1:4**j) = outer_prod_l_(mask_try(1:4**(j-1)), dim_mask_try(:,j))

         end do

         ! See if mask_try falls within mask

         if (COUNT(mask_try .AND. .NOT. mask) == 0) then

            dim_mask = dim_mask_try

            n_match = COUNT(mask_try .AND. mask)

         endif

      case(3)

         ! Cubic case: iterate over all possible combinations of
         ! per-dimension mask, picking the one that produces a
         ! rank-dimensional mask with the most number of points

         do i = 1, 4**rank

            ! Extract bitfields from i
         
            bits_l = IBITS(i-1, rank, rank)
            bits_r = IBITS(i-1, 0, rank)

            ! Accumulate the outer product of per-dimension masks

            mask_try(1) = .TRUE.

            do j = 1, rank

               ! Set up the per-dimension mask

               dim_mask_try(1,j) = BTEST(bits_l, j-1) 
               dim_mask_try(2:3,j) = .TRUE.
               dim_mask_try(4,j) = BTEST(bits_r, j-1)

               ! Accumulate it

               mask_try(1:4**j) = outer_prod_l_(mask_try(1:4**(j-1)), dim_mask_try(:,j))

            end do

            ! See if mask_try falls within mask

            if (COUNT(mask_try .AND. .NOT. mask) == 0) then

               ! If so, compare it to the current best match

               n_match_try = COUNT(mask_try .AND. mask)

               if (n_match_try > n_match) then

                  ! Replace best match

                  dim_mask = dim_mask_try

                  n_match = n_match_try

               elseif (n_match_try == n_match) then

                  ! Merge with best match

                  dim_mask = dim_mask .AND. dim_mask_try

               endif

            endif

         end do

      case default

         if (PRESENT(stat)) then
            stat = STAT_INVALID_ARGUMENT
            return
         else
            @:ABORT('invalid argument')
         end if

      end select
      
      ! Check that matches were found

      if (n_match == 0) then

         if (PRESENT(stat)) then
            stat = STAT_UNAVAILABLE_DATA
            return
         else
            @:ABORT('unavailable data')
         endif

      endif

      ! Now build the kernel matrix row-by-row

      row_loop: do i = 1, 4**rank
      
         ! Extract bitfields from i
         
         bits_l = IBITS(i-1, rank, rank)
         bits_r = IBITS(i-1, 0, rank)

         ! Accumulate the outer product of per-dimension kernels

         K_D(i,1) = 1._RD

         dim_loop : do j = 1, rank

            ! Set up the per-dimension kernel

            if (BTEST(bits_r, j-1)) then

               if (BTEST(bits_l, j-1)) then

                  dim_K(1) = 0
                  dim_K(2:4) = deriv_kernel_(vdelta(2:3,j), dim_mask(2:4,j))

               else

                  dim_K = [0,0,1,0]

               endif

            else

               if (BTEST(bits_l, j-1)) then

                  dim_K(1:3) = deriv_kernel_(vdelta(1:2,j), dim_mask(1:3,j))
                  dim_K(4) = 0

               else

                  dim_K = [0,1,0,0]

               endif

            end if

            ! Accumulate it

            K_D(i,1:4**j) = outer_prod_r_(K_D(i,1:4**(j-1)), dim_K)

         end do dim_loop

      end do row_loop

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   contains

      function deriv_kernel_(delta, mask) result(K)

         real(RD), intent(in) :: delta(:)
         logical, intent(in)  :: mask(:)
         real(RD)             :: K(3)

         real(RD) :: u
         real(RD) :: K_a(3)
         real(RD) :: K_b(3)

         @:CHECK_BOUNDS(SIZE(delta), 2)
         @:CHECK_BOUNDS(SIZE(mask), 3)

         ! Set up a 3-point kernel for derivative evaluation

         if (mask(1) .AND. mask(3)) then

            u = delta(1)/(delta(1) + delta(2))

            K_a = [-1,1,0]/delta(1)
            K_b = [0,-1,1]/delta(2)

            K = (1._RD-u)*K_a + u*K_b

         elseif (mask(1)) then

            K = [-1,1,0]/delta(1)

         elseif (mask(3)) then

            K = [0,-1,1]/delta(2)

         else

            @:ABORT('invalid mask')

         endif

         ! Finish

         return

      end function deriv_kernel_

      !****

      function outer_prod_r_(A, B) result (AB)

         real(RD), intent(in) :: A(:)
         real(RD), intent(in) :: B(:)
         real(RD)             :: AB(SIZE(A)*SIZE(B))

         integer :: i
         integer :: n
         integer :: m

         ! Calculate the outer product between (flat) real arrays A and B

         n = SIZE(A)
         m = SIZE(B)

         do i = 1, m

            AB(n*(i-1)+1:n*i) = A*B(i)

         end do

         ! Finish

         return

      end function outer_prod_r_

      !****

      function outer_prod_l_(A, B) result (AB)

         logical, intent(in) :: A(:)
         logical, intent(in) :: B(:)
         logical             :: AB(SIZE(A)*SIZE(B))

         integer :: i
         integer :: n
         integer :: m

         ! Calculate the outer product between (flat) logical arrays A and B

         n = SIZE(A)
         m = SIZE(B)

         do i = 1, m

            AB(n*(i-1)+1:n*i) = A .AND. B(i)

         end do

         ! Finish

         return

      end function outer_prod_l_

   end subroutine eval_deriv_kernel

   !****

   subroutine eval_interp_kernel(rank, vu, vdelta, K_I, vderiv)

      integer, intent(in)           :: rank
      real(RD), intent(in)          :: vu(:)
      real(RD), intent(in)          :: vdelta(:,:)
      real(RD), intent(out)         :: K_I(:)
      logical, intent(in), optional :: vderiv(:)

      logical  :: vderiv_(rank)
      integer  :: i
      integer  :: p
      integer  :: d
      integer  :: j

      @:CHECK_BOUNDS(SIZE(vu), rank)

      @:CHECK_BOUNDS(SIZE(vdelta, 1), 3)
      @:CHECK_BOUNDS(SIZE(vdelta, 2), rank)

      @:CHECK_BOUNDS(SIZE(K_I), 4**rank)

      if (PRESENT(vderiv)) then
         @:CHECK_BOUNDS(SIZE(vderiv), rank)
         vderiv_ = vderiv
      else
         vderiv_ = .FALSE.
      endif

      ! Evaluate the kernel vector K_I for interpolating. Given a vector
      ! D of function values and derivatives on a 2**rank-point stencil,
      ! DOT_PRODUCT(K_I, D) gives the interpolant at normalized position
      ! vu in the stencil.

      do i = 1, 4**rank

         ! Extract position and derivative indices

         p = IBITS(i-1, 0, rank)
         d = IBITS(i-1, rank, rank)

         ! Evaluate the product of the basis function for each
         ! dimension

         K_I(i) = 1._RD
         
         dim_loop : do j = 1, rank

            if (BTEST(p, j-1)) then
               if (BTEST(d, j-1)) then
                  K_I(i) = K_I(i)*h_11(vu(j), vderiv_(j))*vdelta(2,j)
               else
                  K_I(i) = K_I(i)*h_01(vu(j), vderiv_(j))
               endif
            else
               if (BTEST(d, j-1)) then
                  K_I(i) = K_I(i)*h_10(vu(j), vderiv_(j))*vdelta(2,j)
               else
                  K_I(i) = K_I(i)*h_00(vu(j), vderiv_(j))
               endif
            endif
               
            if (vderiv_(j)) K_I(i) = K_I(i)/vdelta(2,j)

         end do dim_loop

      end do
      
      ! Finish

   contains

      function h_00(u, deriv)

         real(RD), intent(in) :: u
         logical, intent(in)  :: deriv
         real(RD)             :: h_00

         ! Set up the h_00 hermite basis function

         if (deriv) then
            h_00 = 6*u**2 - 6*u
         else
            h_00 = 2*u**3 - 3*u**2 + 1
         endif

         ! Finish

         return

      end function h_00
         
      function h_01(u, deriv)

         real(RD), intent(in) :: u
         logical, intent(in)  :: deriv
         real(RD)             :: h_01

         ! Set up the h_01 hermite basis function

         if (deriv) then
            h_01 = -6*u**2 + 6*u
         else
            h_01 = -2*u**3 + 3*u**2
         endif

         ! Finish

         return

      end function h_01
         
      function h_10(u, deriv)

         real(RD), intent(in) :: u
         logical, intent(in)  :: deriv
         real(RD)             :: h_10

         ! Set up the h_10 hermite basis function

         if (deriv) then
            h_10 = 3*u**2 - 4*u + 1
         else
            h_10 = u**3 - 2*u**2 + u
         endif
         
         ! Finish

         return

      end function h_10
         
      function h_11(u, deriv)

         real(RD), intent(in) :: u
         logical, intent(in)  :: deriv
         real(RD)             :: h_11

         ! Set up the h_11 hermite basis function

         if (deriv) then
            h_11 = 3*u**2 - 2*u
         else
            h_11 = u**3 - u**2
         endif

         ! Finish

         return

      end function h_11

   end subroutine eval_interp_kernel

end module ninterp_m
