! Module  : fit_m
! Purpose : data fitting routines
!
! 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 fit_m

  ! Uses

  use forum_m

  use f95_lapack

  use ISO_FORTRAN_ENV

  ! No implicit typing

  implicit none

  ! Interfaces

  interface gls_fit
     module procedure gls_fit_1_
     module procedure gls_fit_2_
  end interface gls_fit

  ! Access specifiers

  private

  public :: gls_fit

contains

  subroutine gls_fit_1_(x, y, f, a)

    real(RD), intent(in) :: x(:)
    real(RD), intent(in) :: y(:)
    interface
       function f (x, n)
         use forum_m
         real(RD), intent(in) :: x
         integer, intent(in)  :: n
         real(RD)             :: f(n)
       end function f
    end interface
    real(RD), intent(out) :: a(:)

    real(RD), allocatable :: y_2(:,:)
    real(RD), allocatable :: a_2(:,:)

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

    ! Determine the coefficients a_{i} in the fit to the y(x) data
    ! using the function f(x) = sum_{i} a_{i} * f_{i}(x)

    allocate(y_2(SIZE(y), 1))
    allocate(a_2(SIZE(a), 1))

    y_2(:,1) = y

    call gls_fit(x, y_2, f, a_2)

    a = a_2(:,1)

    ! Finish

    return

 end subroutine gls_fit_1_

  !****
  
  subroutine gls_fit_2_(x, y, f, a)

    real(RD), intent(in) :: x(:)
    real(RD), intent(in) :: y(:,:)
    interface
       function f (x, n)
         use forum_m
         real(RD), intent(in) :: x
         integer, intent(in)  :: n
         real(RD)             :: f(n)
       end function f
    end interface
    real(RD), intent(out) :: a(:,:)

    integer               :: n_a
    integer               :: n_x
    integer               :: i
    real(RD), allocatable :: M(:,:)
    real(RD), allocatable :: B(:,:)

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

    ! Determine the coefficients a_{i} in the fit to the y(x,:) data
    !  using the function f(x) = sum_{i} a_{i} * f_{i}(x)

    ! Set up the desgin matrix

    n_a = SIZE(a, 1)
    n_x = SIZE(x)

    allocate(M(n_x,n_a))

    x_loop : do i = 1, n_x
       M(i,:) = f(x(i), n_a)
    end do x_loop

    ! Set up the rhs matrix

    allocate(B(MAX(n_x, n_a), SIZE(y, 2)))

    B(:n_x,:) = y
    B(n_x+1:,:) = 0._RD

    ! Do the fit

    call LA_GELSS(M, B)

    ! Extract the results

    a = B(:n_a,:)

    ! Finish

 end subroutine gls_fit_2_

end module fit_m

