!
!  general interpolation routines by Numerical recipes
!
!  Copyright © 1997-2011 F.Hroch (hroch@physics.muni.cz)
!
!  This file is part of Munipack.
!
!  Munipack 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, either version 3 of the License, or
!  (at your option) any later version.
!  
!  Munipack 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 Munipack.  If not, see <http://www.gnu.org/licenses/>.
!


module interpol

  implicit none
  integer, parameter, private :: sp = selected_real_kind(6)
  integer, private :: errno
  logical, parameter, private :: stopme = .false.

! test for spline boundary values
  real(sp), parameter, private :: hbound = 0.9*huge(1.0_sp)

! work arrays for spline
  real(sp), dimension(:), allocatable, private :: x1a, x2a
  real(sp), dimension(:,:), allocatable, private :: y2a

contains

  function interpol_errno()

    implicit none
    integer :: interpol_errno

    interpol_errno = errno

  end function interpol_errno

  subroutine nrerror(no,string)

    implicit none
    character(len=*), intent(in) :: string
    integer, intent(in) :: no

    errno = no
    if( stopme ) then
       write(*,*) trim(string)
       stop 'Program terminated by nrerror.'
    endif

  end subroutine nrerror

  function assert_eq(name,n1,n2,n3,n4,n5,n6,n7)

    implicit none

    integer, intent(in) :: n1,n2
    character(len=*) :: name
    integer, optional,intent(in) :: n3,n4,n5,n6,n7
    integer :: assert_eq
    logical :: eq

    eq = n1 == n2
    if( present(n3) ) eq = eq .and. n1 == n3
    if( present(n4) ) eq = eq .and. n1 == n4
    if( present(n5) ) eq = eq .and. n1 == n5
    if( present(n6) ) eq = eq .and. n1 == n6
    if( present(n7) ) eq = eq .and. n1 == n7
    
    if( eq ) then
       assert_eq = n1
       errno = 0
    else
       errno = 1
       if( stopme ) then
          write(*,*) 'Dimensions non-equal in ',trim(name),'.'
          stop 'assert_eq'
       endif
    end if

  end function assert_eq


  subroutine spline(x,y,yp1,ypn,y2)

    implicit none

    real(sp), dimension(:), intent(in) :: x,y
    real(sp), intent(in) :: yp1,ypn
    real(sp), dimension(:), intent(out) :: y2
    
    integer :: n
    real(sp), dimension(size(x)) :: a,b,c,r

    n = assert_eq('spline',size(x),size(y),size(y2))

    c(1:n-1) = x(2:n) - x(1:n-1)
    r(1:n-1) = 6.0_sp*(y(2:n) - y(1:n-1))/c(1:n-1)
    r(2:n-1) = r(2:n-1) - r(1:n-2)
    a(2:n-1) = c(1:n-2)
    b(2:n-1) = 2.0_sp*(c(2:n-1) + a(2:n-1))
    b(1) = 1.0
    b(n) = 1.0
    if( yp1 > 0.99*hbound ) then
       r(1) = 0.0
       c(1) = 0.0
    else
       r(1) = (3.0_sp/(x(2) - x(1)))*((y(2) - y(1))/(x(2) - x(1)) - yp1)
       c(1) = 0.5_sp
    end if
    if( ypn > 0.99*hbound )then
       r(n) = 0.0
       a(n) = 0.0
    else
       r(n) = (-3.0_sp/(x(n) - x(n-1)))*((y(n) - y(n-1))/(x(n) - x(n-1)) - ypn)
       a(n) = 0.5
    end if
    call tridag_ser(a(2:n),b(1:n),c(1:n-1),r(1:n),y2(1:n))

  end subroutine spline

  function splint(xa,ya,y2a,x)

    implicit none
    
    real(sp), dimension(:), intent(in) :: xa,ya,y2a
    real(sp), intent(in) :: x
    real(sp) :: splint

    integer :: khi,klo,n
    real(sp) :: a,b,h

    n = assert_eq('splint',size(xa),size(ya),size(y2a))
    
    klo = max(min(locate(xa,x),n-1),1)
    khi = klo + 1
    h = xa(khi) - xa(klo)

    if( h == 0.0 ) call nrerror(4,'Bad xa input in splint.')
    a = (xa(khi) - x)/h
    b = (x - xa(klo))/h
    splint = a*ya(klo) + b*ya(khi) + ((a**3 - a)*y2a(klo) + &
         (b**3 - b)*y2a(khi))*(h**2)/6.0_sp

  end function splint


  subroutine tridag_ser(a,b,c,r,u)

    implicit none

    real(sp), dimension(:), intent(in) :: a,b,c,r
    real(sp), dimension(:), intent(out) :: u

    real(sp), dimension(size(b)) :: gam
    integer :: n,j
    real(sp) :: bet

    n = assert_eq('tridag_ser',size(a)+1,size(b),size(c)+1,size(r),size(u))

    bet = b(1)
    if( bet == 0.0 ) call nrerror(2,'tridag_ser: Error at code stage 1')
    u(1) = r(1)/bet
    do j = 2,n
       gam(j) = c(j-1)/bet
       bet = b(j) - a(j-1)*gam(j)
       if( bet == 0.0 ) call nrerror(3,'tridag_ser: Error at code stage 2')
       u(j) = (r(j) - a(j-1)*u(j-1))/bet
    enddo
    do j = n-1,1,-1
       u(j) = u(j) - gam(j+1)*u(j+1)
    enddo
    
  end subroutine tridag_ser

  function locate(xx,x)

    implicit none

    real(sp), dimension(:), intent(in) :: xx
    real(sp), intent(in) :: x
    integer :: locate

    integer :: n,jl,jm,ju
    logical :: ascnd

    n = size(xx)
    ascnd = (xx(n) >= xx(1))
    jl = 0
    ju = n + 1

    do
       if( ju - jl <= 1 ) exit
       jm = (ju + jl)/2
       if( ascnd .eqv. (x >= xx(jm)) )then
          jl = jm
       else
          ju = jm
       end if
    end do

    if( x == xx(1) )then
       locate = 1
    else if(x == xx(n) )then
       locate = n - 1
    else
       locate = jl
    end if

  end function locate

  subroutine bcucof(y,y1,y2,y12,d1,d2,c)

    implicit none

    real(sp), intent(in) :: d1,d2
    real(sp), dimension(4), intent(in) :: y,y1,y2,y12
    real(sp), dimension(4,4), intent(out) :: c

    real(sp), dimension(16) :: x

    real(sp), parameter, dimension(16,16) :: wt = reshape( (/ &
     1,0,-3,2,0,0,0,0,-3,0,9,-6,2,0,-6,4,0,0,0,0,0,0,0,0,3,0,-9,6,-2,0,6,-4, &
     0,0,0,0,0,0,0,0,0,0,9,-6,0,0,-6,4,0,0,3,-2,0,0,0,0,0,0,-9,6,0,0,6,-4, &  
     0,0,0,0,1,0,-3,2,-2,0,6,-4,1,0,-3,2,0,0,0,0,0,0,0,0,-1,0,3,-2,1,0,-3,2,& 
     0,0,0,0,0,0,0,0,0,0,-3,2,0,0,3,-2,0,0,0,0,0,0,3,-2,0,0,-6,4,0,0,3,-2,&
     0,1,-2,1,0,0,0,0,0,-3,6,-3,0,2,-4,2,0,0,0,0,0,0,0,0,0,3,-6,3,0,-2,4,-2,&
     0,0,0,0,0,0,0,0,0,0,-3,3,0,0,2,-2,0,0,-1,1,0,0,0,0,0,0,3,-3,0,0,-2,2,&
     0,0,0,0,0,1,-2,1,0,-2,4,-2,0,1,-2,1,0,0,0,0,0,0,0,0,0,-1,2,-1,0,1,-2,1, &
     0,0,0,0,0,0,0,0,0,0,1,-1,0,0,-1,1,0,0,0,0,0,0,-1,1,0,0,2,-2,0,0,-1,1 /),&
     (/16,16/) )

    x(1:4) = y
    x(5:8) = y1*d1
    x(9:12) = y2*d2
    x(13:16) = y12*d1*d2
    x = matmul(wt,x)
    c = reshape(x,(/4,4/),order=(/2,1/))
    
  end subroutine bcucof

  subroutine bcuint(y,y1,y2,y12,x1l,x1u,x2l,x2u,x1,x2,ansy,ansy1,ansy2)

    implicit none

    real(sp), dimension(4), intent(in) :: y,y1,y2,y12
    real(sp), intent(in) :: x1l,x1u,x2l,x2u,x1,x2
    real(sp), intent(out) :: ansy,ansy1,ansy2

    integer :: i
    real(sp) :: t,u
    real(sp), dimension(4,4) :: c

    call bcucof(y,y1,y2,y12,x1u-x1l,x2u-x2l,c)

    if( x1u == x1l .or. x2u == x2l ) call &
         nrerror(6,'bcuint: problem with input values - boundary pair equal?')

    t = (x1 - x1l)/(x1u - x1l)
    u = (x2 - x2l)/(x2u - x2l)
    ansy = 0.0
    ansy1 = 0.0
    ansy2 = 0.0
    do i = 4,1,-1
       ansy = t*ansy + ((c(i,4)*u + c(i,3))*u + c(i,2))*u + c(i,1)
       ansy1 = u*ansy1 + (3.0*c(4,i)*t + 2.0*c(3,i))*t + c(2,i)
       ansy2 = t*ansy2 + (3.0*c(i,4)*u + 2.0*c(i,3))*u + c(i,2)
    end do
    ansy1 = ansy1/(x1u - x1l)
    ansy2 = ansy2/(x2u - x2l)

  end subroutine bcuint

  subroutine spline2(x1a,x2a,ya,y2a)

    implicit none

    real(sp), dimension(:), intent(in) :: x1a,x2a
    real(sp), dimension(:,:), intent(in) :: ya
    real(sp), dimension(:,:), intent(out) :: y2a

    integer :: j,m,ndum
    
    m = assert_eq('spline2: m',size(x1a),size(ya,1),size(y2a,1))
    ndum = assert_eq('spline2: ndum',size(x2a),size(ya,2),size(y2a,2))
    do j = 1,m
       call spline(x2a,ya(j,:),hbound,hbound,y2a(j,:))
    end do
    
  end subroutine spline2

  function splint2(x1a,x2a,ya,y2a,x1,x2)

    implicit none
    
    real(sp), dimension(:), intent(in) :: x1a,x2a
    real(sp), dimension(:,:), intent(in) :: ya,y2a
    real(sp), intent(in) :: x1,x2
    real(sp) :: splint2

    integer :: j,m,ndum
    real(sp), dimension(size(x1a)) :: yytmp, y2tmp
    
    m = assert_eq('splint2: m',size(x1a),size(ya,1),size(y2a,1))
    ndum = assert_eq('splint2: ndum',size(x2a),size(ya,2),size(y2a,2))

    do j = 1,m
       yytmp(j) = splint(x2a,ya(j,:),y2a(j,:),x2)
    end do
    
    call spline(x1a,yytmp,hbound,hbound,y2tmp)
    splint2 = splint(x1a,yytmp,y2tmp,x1)

  end function splint2

  function bilinear(x,y,array) result(fun)

    real(sp), intent(in) :: x,y
    real(sp), dimension(:,:), intent(in) :: array
    real(sp) :: fun
    integer :: i,j
    real(sp) :: u,v

    i = int(x)
    j = int(y)

    if( (lbound(array,1) <= i .and. i < ubound(array,1)) .and. &
        (lbound(array,2) <= j .and. j < ubound(array,2)) ) then
         
       u = (x - i)/(i+1 - i)
       v = (y - j)/(j+1 - j)

       fun = (1.0 - u)*(1.0 - v)*array(i,j) + u*v*array(i+1,j+1) &
            + u*(1.0 - v)*array(i+1,j) + (1.0 - u)*v*array(i,j+1)
    
       
    else
       fun = 0.0
    end if

  end function bilinear


  function bicubic(x,y,array)

    implicit none

    real(sp), intent(in) :: x,y
    real(sp), dimension(:,:), intent(in) :: array
    real(sp) :: bicubic

    integer :: i,j,nx,ny
    real(sp) :: d, z, z1, z2, x1l, x1u, x2l, x2u
    real(sp), dimension(4) :: f, f1, f2, f12

    nx = size(array,1)
    ny = size(array,2)
    i = int(x)
    j = int(y)

    d = 1.0  ! coordinate step

    if( 1 < i .and. i < nx - 1 .and. 1 < j .and. j < ny - 1 ) then

       f(1) = array(i,j)
       f(2) = array(i+1,j)
       f(3) = array(i+1,j+1)
       f(4) = array(i,j+1)

       f1(1) = (array(i+1,j)   - array(i-1,j))  /(2.0*d)
       f1(2) = (array(i+2,j)   - array(i,j))    /(2.0*d)
       f1(3) = (array(i+2,j+1) - array(i,j+1))  /(2.0*d)
       f1(4) = (array(i+1,j+1) - array(i-1,j+1))/(2.0*d)

       f2(1) = (array(i,j+1)   - array(i,j-1))  /(2.0*d)
       f2(2) = (array(i+1,j+1) - array(i+1,j-1))/(2.0*d)
       f2(3) = (array(i+1,j+2) - array(i+1,j))  /(2.0*d)
       f2(4) = (array(i,j+2)   - array(i,j))    /(2.0*d)

       f12(1) = (array(i+1,j+1)-array(i+1,j-1)-array(i-1,j+1)+array(i-1,j-1))/(2.0*d)**2
       f12(2) = (array(i+2,j+1)-array(i+2,j-1)-array(i,j+1)+array(i,j-1))/(2.0*d)**2
       f12(3) = (array(i+2,j+2)-array(i+2,j)-array(i,j+2)+array(i,j))/(2.0*d)**2
       f12(4) = (array(i+1,j+2)-array(i+1,j)-array(i-1,j+2)+array(i-1,j))/(2.0*d)**2

       x1l = i
       x1u = x1l + d
       x2l = j
       x2u = x2l + d
       
       call bcuint(f,f1,f2,f12,x1l,x1u,x2l,x2u,x,y,z,z1,z2)
       bicubic = z

    else
       
       bicubic = 0

    end if

  end function bicubic

  function bspline(x,y,array)

    implicit none

    real(sp), intent(in) :: x,y
    real(sp), dimension(:,:), intent(in) :: array
    real(sp) :: bspline

    call bspline_init(array)

    bspline = bspline_func(x,y,array)

    call bspline_shutdown

  end function bspline

  subroutine bspline_init(array)

    implicit none

    real(sp), dimension(:,:), intent(in) :: array

    integer :: i

    allocate(x1a(size(array,1)))
    allocate(x2a(size(array,2)))
    allocate(y2a(size(array,1),size(array,2)))

    do i = 1, size(x1a)
       x1a(i) = i
    enddo
    do i = 1, size(x2a)
       x2a(i) = i
    enddo

    call spline2(x1a,x2a,array,y2a)

  end subroutine bspline_init

  subroutine bspline_shutdown

    implicit none

    deallocate(x1a)
    deallocate(x2a)
    deallocate(y2a)

  end subroutine bspline_shutdown

  function bspline_func(x,y,array)

    implicit none

    real(sp), intent(in) :: x,y
    real(sp), dimension(:,:), intent(in) :: array
    real(sp) :: bspline_func

    bspline_func = splint2(x1a,x2a,array,y2a,x,y)

  end function bspline_func

end module interpol


