module splines
  ! sets working precision (wp)
  use set_precision, only: wp
  ! use tridiagonal solvers
  use tridiag
  implicit none

  ! spline data structure
  type :: spline
     ! length is current length
     ! max_len is allocated length of arrays
     integer length, max_len
     ! For spline function f
     ! f(xlist(i)) = ylist(i), Mlist(i) = f''(xlist(i))
     real(kind=wp), pointer :: &
          xlist(:), ylist(:), Mlist(:)
  end type spline
  
contains
  ! Spline functions
  ! spline_get(s,len) -- creates spline s of length len
  !  -- returns .true. if successful, .false. if not
  logical function spline_get(s,len)
    type(spline), intent(out) :: s
    integer, intent(in) :: len
    integer :: err
    s%length = 0
    s%max_len = 0
    allocate(s%xlist(len), s%ylist(len), s%Mlist(len), &
         STAT=err)
    if ( err /= 0 ) then
       deallocate(s%xlist, s%ylist, s%Mlist)
       spline_get = .false.
    else
       s%length = len
       s%max_len = len
       spline_get = .true.
    end if
  end function spline_get
  
  ! spline_free(s) -- deallocates all memory associated 
  !   with s and resets length and max_len.
  subroutine spline_free(s)
    type(spline) :: s
    deallocate(s%xlist, s%ylist, s%Mlist)
    s%length = 0
    s%max_len = 0
  end subroutine spline_free
  
  ! spline_resize(s,new_len) -- resize s to length new_len
  !  -- re-allocate if necessary, possibly destroying old values
  logical function spline_resize(s,new_len)
    type(spline) :: s
    integer err, new_len
  
    spline_resize = .true.
    if ( new_len > s%length ) then
       ! deallocate and allocate
       deallocate(s%xlist, s%ylist, s%Mlist)
       s%length = 0
       allocate(s%xlist(new_len), s%ylist(new_len), &
            s%Mlist(new_len), STAT = err)
       if ( err /= 0 ) then
          spline_resize = .false.
          return
       end if
       s%length = new_len
       s%max_len = new_len
    else if ( new_len >= 0 ) then
       s%length = new_len
    end if
  end function spline_resize
  
  ! spline_set(s,x,y,len) -- sets xlist & ylist values in s
  !        to x & y respectively
  !  -- resizes s to ensure it has length = len
  !  -- x must be ordered on entry: x(i+1) > x(i) for all i
  !      but this is not checked by spline_set
  logical function spline_set(s,x,y,len)
    type(spline) :: s
    real(kind=wp) x(:), y(:)
    integer i, len
    spline_set = .true.
    if ( .not. spline_resize(s,len) ) then
       spline_set = .false.
       return
    end if
    ! check that x(:) is an increasing array
    do i = 1, len-1
       if ( x(i) >= x(i+1) ) then
          spline_set = .false.
          return
       end if
    end do
    do i = 1, len
       s%xlist(i) = x(i)
       s%ylist(i) = y(i)
    end do
  end function spline_set
  
  ! spline_print(s,output) -- prints representation of spline 
  !   on unit "output"
  subroutine spline_print(s,output)
    type(spline) :: s
    integer i, output
    write (output,*) 'Spline:', s%length
    write (output,*) 'xlist:'
    write (output,'(3G22.15)') (s%xlist(i), i = 1, s%length)
    write (output,*) 'ylist:'
    write (output,'(3G22.15)') (s%ylist(i), i = 1, s%length)
    write (output,*) 'Mlist:'
    write (output,'(3G22.15)') (s%Mlist(i), i = 1, s%length)
  end subroutine spline_print
  
  ! spline_read(s,input) -- reads representation of spline 
  !   on unit "input" (reads output from spline_print)
  logical function spline_read(s,input)
    type(spline) :: s
    integer :: err, input, i, len
    character(80) :: string
  
    spline_read = .true.
    read(input,*, IOSTAT=err) string, len  ! read 'String: <length>'
    if ( err /= 0 .or. .not. spline_resize(s,len) ) then
       spline_read = .false.
       return
    end if
    read(input,*) string       ! read 'xlist:'
    read(input,*, IOSTAT=err) (s%xlist(i), i = 1, len)
    read(input,*) string       ! read 'ylist:'
    if ( err == 0 ) then
       read(input,*, IOSTAT=err) (s%ylist(i), i = 1, len)
    end if
    read(input,*) string       ! read 'Mlist:'
    if ( err == 0 ) then
       read(input,*, IOSTAT=err) (s%Mlist(i), i = 1, len)
    end if
    if ( err /= 0 ) spline_read = .false.
  end function spline_read
  
  ! spline_eval(s,t,iflag) -- evaluate spline s at t: returns s(t)
  !   -- extrapolates if t outside interval 
  !    [xlist(1),xlist(length)], but then sets iflag > 0
  !   -- if fatal error (e.g., division by zero), set iflag < 0
  !   -- otherwise, set iflag == 0 on exit
  function spline_eval(s,t,iflag) result(val)
    type(spline), intent(in) :: s
    integer, intent(out) :: iflag
    real(kind=wp) :: t, val, h, diff1, diff2
    integer i, lo, mid, hi
  
    if ( s%length == 1 ) then
       val = s%ylist(1)
       iflag = 1 ! extrapolation !
       return
    else if ( s%length == 0 ) then
       val = 0
       iflag = -1 ! no data !
       return
    end if
  
    ! binary search to find i: xlist(i) <= t <= xlist(i+1),
    ! or i = 1 if t < xlist(1), 
    ! or i = length-1 if xlist(length) < t
    lo = 1
    hi = s%length
    if ( t < s%xlist(1) ) then
       i = 1
       iflag = 1 ! extrapolation !
    else if ( t > s%xlist(s%length) ) then
       i = s%length-1
       iflag = 1 ! extrapolation !
    else
       iflag = 0 ! normal case !
       ! binary search
       do while ( hi > lo + 1 )
          ! Loop invariant: xlist(lo) <= t <= xlist(hi)
          mid = (lo + hi)/2
          if ( t <= s%xlist(mid) ) then
             hi = mid
          else
             lo = mid
          end if
       end do
       i = lo
    end if
  
    ! Use interpolation formula on [xlist(i),xlist(i+1)]
    h = s%xlist(i+1) - s%xlist(i)
    if ( h == 0 ) then
       iflag = -1 ! fatal error !
       return
    end if
    diff1 = t - s%xlist(i)
    diff2 = s%xlist(i+1) - t
    val = diff1*(s%ylist(i+1)/h - s%Mlist(i+1)*h/6.0 &
         + diff1*diff1*s%Mlist(i+1)/(6.0*h)) &
         + diff2*(s%ylist(i )/h - s%Mlist(i  )*h/6.0 &
         + diff2*diff2*s%Mlist(i  )/(6.0*h))
  end function spline_eval
  
  ! spline_peval(s,t,iflag) -- evaluate spline s at t: returns s(t)
  !   -- assumes spline is periodic
  !   -- if fatal error, set iflag < 0
  function spline_peval(s,t,iflag) result(val)
    type(spline) :: s
    integer, intent(out) :: iflag
    real(kind=wp) :: t, val, t_reduced, p
  
    iflag = 0
    if ( s%length == 1 ) then
       val = s%ylist(1)
       return
    else if ( s%length == 0 ) then
       iflag = -1
       val = 0
       return
    end if
  
    p = s%xlist(s%length) - s%xlist(1)
    t_reduced = s%xlist(1) + modulo(t-s%xlist(1), p)
    val = spline_eval(s,t_reduced,iflag)
    
  
    ! ignore non-fatal errors
    if ( iflag > 0 ) iflag = 0
  end function spline_peval
  
  
  ! spline_basic(s,mat,rhs,iflag) -- constructs basic linear system 
  !    for splines for computing Mlist
  ! This is common to spline_clamped, spline_natural, 
  ! and spline_not_knot; sets iflag < 0 if fatal error
  subroutine spline_basic(s,mat,rhs,iflag)
    type(spline), intent(in) :: s
    type(tdg) :: mat
    integer, intent(out) :: iflag
    real(kind=wp) rhs(:), h(s%length-1)
    integer i, n
  
    iflag = 0
    n = s%length
    if ( .not. tdg_resize(mat,n) ) then
       iflag = -1
       return
    end if
  
    ! Set up h array
    do i = 1, n - 1
       h(i) = s%xlist(i+1) - s%xlist(i)
       if ( h(i) == 0 ) then
          iflag = -1
          return
       end if
    end do
    ! Set up right-hand side (rhs) and matrix entries
    do i = 2, n - 1
       rhs(i) = (s%ylist(i+1)-s%ylist(i))/h(i) &
            - (s%ylist(i)-s%ylist(i-1))/h(i-1)
       mat%b(i-1) = h(i-1)/6.
       mat%a(i)   = (h(i-1)+h(i))/3
       mat%c(i)   = h(i)/6
    end do
  end subroutine spline_basic
  
  
  ! spline_clamped(s,dy0,dy1,iflag) -- constructs clamped spline f(.):
  !    f(s%xlist(i)) = s%ylist(i), i = 1, ..., s%length
  !    f'(s%xlist(1)) = dy0, f'(s%xlist(s%length)) = dy1
  subroutine spline_clamped(s,dy0,dy1,iflag)
    type(tdg) :: mat
    type(spline) :: s
    integer, intent(out) :: iflag
    real(kind=wp) dy0, dy1, h_1, h_n_1
    integer i, n
  
    iflag = 0
    n = s%length
    if ( .not. tdg_get(mat,n) ) then
       iflag = -1
       return
    end if
    call spline_basic(s,mat,s%Mlist,iflag)
    if ( iflag == 0 ) then
       ! Now we must set the first and last rows
       h_1 = s%xlist(2) - s%xlist(1)
       h_n_1 = s%xlist(n) - s%xlist(n-1)
       mat%a(1) = h_1/3
       mat%c(1) = h_1/6
       s%Mlist(1) = (s%ylist(2)-s%ylist(1))/h_1 - dy0
       mat%b(n-1) = h_n_1/6
       mat%a(n)   = h_n_1/3
       s%Mlist(n) = dy1 - (s%ylist(n)-s%ylist(n-1))/h_n_1
    end if
    
    
    ! call tdg_print(mat,6)
    ! print *, (s%Mlist(i), i = 1, n)
  
    ! With the linear system set up, we solve for Mlist
    if ( iflag == 0 ) then
       call tdg_lu(mat,iflag)
    end if
    if ( iflag == 0 ) then
       call tdg_lusolve(mat,s%Mlist,iflag)
    end if
  
    call tdg_free(mat)
  end subroutine spline_clamped
  
  
  ! spline_natural(s,iflag) -- constructs natural spline f(.):
  !    f(s%xlist(i)) = s%ylist(i), i = 1, ..., s%length
  !    f''(s%xlist(1)) = f''(s%xlist(s%length)) = 0
  subroutine spline_natural(s,iflag)
    type(tdg) :: mat
    type(spline) :: s
    integer, intent(out) :: iflag
    integer :: i, n
  
    n = s%length
    if ( .not. tdg_get(mat,n) ) then
       iflag = -1
       return 
    end if
    call spline_basic(s,mat,s%Mlist,iflag)
    if ( iflag == 0 ) then
       ! Now we must set the first and last rows
       mat%a(1) = 1
       mat%c(1) = 0
       s%Mlist(1) = 0
       mat%b(n-1) = 0
       mat%a(n)   = 1
       s%Mlist(n) = 0
    end if
  
    ! call tdg_print(mat,6)
    ! print *, (s%Mlist(i), i = 1, n)
  
    ! With the linear system set up, we solve for Mlist
    if ( iflag == 0 ) then
       call tdg_lu(mat,iflag)
    end if
    if ( iflag == 0 ) then
       call tdg_lusolve(mat,s%Mlist,iflag)
    end if
    call tdg_free(mat)
  end subroutine spline_natural
  
  ! spline_periodic(s,iflag) -- constructs periodic spline f(.):
  !    f(s%xlist(i)) = s%ylist(i), i = 1, ..., s%length
  ! Period of f is s%xlist(s%length)) - s%xlist(1)
  subroutine spline_periodic(s,iflag)
    type(tdg) :: mat
    type(spline) :: s
    real(kind=wp) :: alpha, b_nm1, h1, h_nm1, h_nm2
    real(kind=wp) :: u(s%length), w(s%length)
    integer i, n
    integer, intent(out) :: iflag
    integer, parameter :: stdout = 6
  
    n = s%length
    if ( .not. tdg_get(mat,n) ) then
       iflag = -1
       return
    end if
    ! Enforce periodicity in ylist
    s%ylist(n) = s%ylist(1)
    ! Set up most of the linear system
    call spline_basic(s,mat,s%Mlist,iflag)
    if ( iflag == 0 ) then
       ! Decouple M_n from the rest of the unknowns
       mat%a(n) = 1
       mat%b(n-1) = 0
       mat%c(n-1) = 0
       s%Mlist(n) = 0 ! Set zero right-hand side for M_n
       ! Set mat to be tridiagonal matrix for application of
       ! Sherman-Morrisson-Woodbury formula as described in text.
       ! This requires modifying rows 1 and n-1.
       h1 = s%xlist(2) - s%xlist(1)
       h_nm1 = s%xlist(n)   - s%xlist(n-1)
       h_nm2 = s%xlist(n-1) - s%xlist(n-2)
       b_nm1 = h_nm1 / 6
       mat%a(1)   = (h1   +h_nm1)/3 - b_nm1
       mat%a(n-1) = (h_nm1+h_nm2)/3 - b_nm1
       mat%c(1)   = mat%b(1)
    end if
    if ( h1 == 0 .or. h_nm1 == 0 .or. h_nm2 == 0 ) then
       iflag = -1
    end if
    if ( iflag == 0 ) then
       ! Set right-hand sides for rows 1 & n-1
       s%Mlist(1)   = (s%ylist(2)-s%ylist(1))/h1 &
            - (s%ylist(1)  -s%ylist(n-1))/h_nm1
       s%Mlist(n-1) = (s%ylist(1)-s%ylist(n-1))/h_nm1 &
            - (s%ylist(n-1)-s%ylist(n-2))/h_nm2
       ! Compute w
       do i = 1, n
          w(i) = 0
       end do
       w(  1) = b_nm1
       w(n-1) = b_nm1
    end if
    if ( iflag == 0 ) then
       call tdg_lu(mat,iflag)
    end if
    if ( iflag == 0 ) then
       call tdg_lusolve(mat,s%Mlist,iflag) ! Now s%Mlist = p
    end if
    if ( iflag == 0 ) then
       call tdg_lusolve(mat,w,iflag)
    end if
    ! Apply Sherman-Morrisson-Woodbury formula
    if ( 1+w(1)+w(n-1) == 0 ) then
       iflag = -1
    end if
    if ( iflag == 0 ) then
       alpha = (s%Mlist(1)+s%Mlist(n-1))/(1+w(1)+w(n-1))
       do i = 1, n-1
          s%Mlist(i) = s%Mlist(i) - alpha*w(i)
       end do
       ! And finally, fix Mlist(n-1)
       s%Mlist(n) = s%Mlist(1)
    end if
  
    call tdg_free(mat)
  end subroutine spline_periodic
  
  
  
  ! Spline tests
  real(kind=wp) function test_spline1(x) result(val)
    real(kind=wp) x, three, six
    three = 3.0
    six = 6.0
    if ( x <= 1.0d0 ) then
       val = -1.0d0 + x + x**3/six
    else if ( x <= 3.0 ) then
       val = 1.0d0/six + 1.5d0*(x-1) + 0.5d0*(x-1)**2 - 0.25d0*(x-1)**3
    else
       val = 3.0d0+1.0d0/six + 0.5d0*(x-3) - (x-3)**2 + 1.0d0/three*(x-3)**3
    end if
  end function test_spline1
  !-----------------------------------------------
  
  real(kind=wp) function test_spline2(x) result(val)
    real(kind=wp) t, x, one, two, three, four, five, six, twelve, alpha
    one = 1.0
    two = 2.0
    three = 3.0
    four = 4.0
    five = 5.0
    six = 6.0
    twelve = 12.0
  
    alpha = one/twelve
    t = modulo(x,six)
    if ( t <= two ) then
       val = alpha*t - t**2 + t**3/four
    else if ( t <= three ) then
       val = 2*alpha-2 + (alpha-1)*(t-2)+(t-2)**2/two
    else if ( t <= five ) then
       val =  3*alpha-2.5+alpha*(t-3)+(t-3)**2/two-(t-3)**3/twelve
    else 
       val = 5*alpha - 7/six + (alpha+1)*(t-5)-(t-5)**3/three
    end if
    val = val + 1
  end function test_spline2
  !-----------------------------------------------
  
  subroutine spline_test1
    type(spline) :: s
    integer i, iflag
    integer, parameter :: stdout = 6
    real(kind=wp) :: x(5), y(5), M(5), points(5), vals(5), diff
  
    print *, 'Spline test 1'
  
    ! Data statements are not used here as this allows for conversion
    ! to different precisions.  No accuracy should be lost in x, y & M
    ! as all these values are exactly represented in single precision
    x(1) = 0.0; x(2) = 1.0; x(3) = 3.0; x(4) = 4.5; x(5) = 5.0
    y(1) = 2.0; y(2) =-1.0; y(3) = 0.5; y(4) = 3.0; y(5) = 1.0
    M(1) = 0.0; M(2) = 1.0; M(3) = 3.0; M(4) = 0.0; M(5) =-1.0
    points(1) = 1.0d0; points(2) = 1.5d0;  points(3) = 1.7d0
    points(4) = 2.3d0; points(5) = 4.6d0
    vals(1)  = -1.0d0; vals(2) =-1.3125d0; vals(3) =-1.3395d0
    vals(4)  = -0.9805d0; vals(5) = 2.608d0
  
    if ( .not. spline_get(s,5) ) then
       print *, 'Error creating spline'
       return
    end if
    if ( .not. spline_set(s,x,y,5) ) then
       print *, 'Error setting spline x, y entries'
       return
    end if
    do i = 1, 5
       s%Mlist(i) = M(i);
    end do
    print *, 'Spline with pre-specified Mlist'
    call spline_print(s,stdout)
  
    print *, 'Errors in evaluations'
    do i = 1, 5
       diff = spline_eval(s,points(i),iflag) - vals(i)
       print *, 'Difference(', i, ') =', diff
    end do
  
    print *, ''
  
    call spline_free(s)
  
  end subroutine spline_test1
  !-----------------------------------------------
  
  
  subroutine spline_test2
    integer, parameter :: stdout = 6, n = 4
    integer i, iflag
    type(spline) :: s
    real(kind=wp) x(n), y(n), M(n), dy0, dy1, max_diff
    real(kind=wp) points(10)
  
    print *, 'Spline test 2'
  
    if ( .not. spline_get(s,n) ) then
       print *, 'Error allocating spline'
       return
    end if
    x(1) = 0.0; x(2) = 1.0; x(3) = 3.0; x(4) = 4.0
    y(1) =-1.0; y(2) =1/6.0d0; y(3) = 3+1/6.0d0; y(4) = 3.0
    M(1) = 0.0; M(2) = 1.0; M(3) = -2.0; M(4) = 0.0
    dy0 = 1; dy1 = -0.50d0
  
    points(1) = 0.3; points(2) = -0.25; points(3) = 1.11
    points(4) = 3.1; points(5) = 2.1;   points(6) = 4.05
    points(7) = 1.6; points(8) = 3.85;  points(9) = 2.5
    points(10) = 3.0
  
    if ( .not. spline_set(s,x,y,n) ) then
       print *, 'Error setting spline x, y entries'
       return
    end if
    print *, 'Spline without Mlist set'
    call spline_print(s,stdout)
    call spline_clamped(s,dy0,dy1,iflag)
    if ( iflag /= 0 ) then
       print *, 'Error computing clamped spline, iflag =', iflag
    end if
    print *, 'Mlist for clamped spline'
    do i = 1, n
       print *, 'Mlist(', i, ') =', s%Mlist(i)
    end do
  
    max_diff = 0
    do i = 1, n
       max_diff = max(max_diff,abs(s%Mlist(i)-M(i)))
    end do
    print *, 'Max error in Mlist =', max_diff
  
    ! Reset Mlist
    do i = 1, n
       s%Mlist(i) = 0
    end do
  
    print *, 'Mlist for natural spline'
    call spline_natural(s,iflag)
    if ( iflag /= 0 ) then
       print *, 'Error computing natural spline: iflag =', iflag
    end if
  
    do i = 1, n
       print *, 'Mlist(', i, ') =', s%Mlist(i)
    end do
  
    max_diff = 0
    do i = 1, n
       max_diff = max(max_diff,abs(s%Mlist(i)-M(i)))
    end do
    print *, 'Max error in Mlist =', max_diff
    print *, ''
  
    print *, 'Recomputation via test_spine1'
    do i = 1, n
       y(i) = test_spline1(x(i))
    end do
    if ( .not. spline_set(s,x,y,n) ) then
       print *, 'Error setting spline x, y entries'
       return
    end if
    dy0 = 1.0
    dy1 = -0.5
    call spline_clamped(s,dy0,dy1,iflag)
    if ( iflag /= 0 ) then
       print *, 'Error computing clamped spline, iflag =', iflag
    end if
    call spline_print(s,stdout)
    max_diff = 0
    do i = 1, 10
       max_diff = max(max_diff,abs( &
            spline_eval(s,points(i),iflag)-test_spline1(points(i))))
    end do
    print *, 'Max error in evaluations =', max_diff
    print *, ''
  
    call spline_free(s)
  
  end subroutine spline_test2
  !-----------------------------------------------
  
  subroutine spline_test3
    type(spline) :: s
    integer, parameter :: stdin = 5, stdout = 6
    !s = spline_get(5)
    if ( .not. spline_get(s,5) ) then
       print *, 'Error allocating spline'
       return
    end if
  
    call spline_print(s,stdout)
  
    !  open(UNIT=10, NAME='splinetest3.txt', ACTION='READ', STATUS='OLD')
    
    call spline_free(s)
  end subroutine spline_test3
  
  !-----------------------------------------------
  
  subroutine spline_test4
    type(spline) :: s
    integer, parameter :: stdin = 5, stdout = 6
    integer :: i, iflag
    real(kind=wp) :: err, max_err
  
    if ( .not. spline_get(s,5) ) then
       print *, 'Error allocating spline'
       return
    end if
  
    s%xlist(1) = 0
    s%xlist(2) = 2
    s%xlist(3) = 3
    s%xlist(4) = 5
    s%xlist(5) = 6
  
    call spline_print(s,stdout)
    print *, 'length =', s%length
    do i = 1, s%length
       s%ylist(i) = test_spline2(s%xlist(i))
    end do
  
    call spline_print(s,stdout)
  
    call spline_periodic(s,iflag)
    if ( iflag /= 0 ) then
       print *, 'Cannot create periodic spline, iflag =', iflag
    end if
  
    call spline_print(s,stdout)
  
    max_err = 0
    max_err = max(max_err,abs(s%Mlist(1)+2))
    max_err = max(max_err,abs(s%MList(2)-1))
    max_err = max(max_err,abs(s%MList(3)-1))
    max_err = max(max_err,abs(s%MList(4)))
    max_err = max(max_err,abs(s%MList(5)+2))
    print *, 'max error =', max_err
  
    call spline_free(s)
  end subroutine spline_test4
  
end module splines