!
!  Robust estimate of parameters of line
!
!  Copyright © 2014-5, 2017 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 robustline

  ! This module provides subroutines for estimation
  ! of parameters (a,b) for a line in the form
  !
  !     y = a + b*x                        (without errors of x,y)
  !
  ! or
  !                a + b*x
  !     y = ------------------------       (with errors)
  !         sqrt(dy**2 + b**2*dx**2)
  !
  !
  ! with optional estimate of theirs statistical errors (da, db).
  !
  ! rline should be called as:
  !
  !  call rline(x,y,a,b,da,db,dx,dy,sigma,reliable,verbose)
  !
  ! on input:
  !   x,y - array of data values to be estimated
  !   dx,dy - array of statistical errors of x,y (optional), dx can be omitted,
  !           if dy is presented.
  !   verbose - verbose output (optional)
  !
  ! on output are estimated:
  !   a,b - parameters of the line
  !   da,db - standard error of the parameters (optional)
  !   sigma - standard deviation (optional)
  !   reliable - indicates reliability of results (optional)

  use rfun

  implicit none

  ! numerical precision of real numbers
  integer, parameter, private :: dbl = selected_real_kind(15)

  ! verbose print?
  logical, private :: xverbose = .false.

  ! print debug messages
  logical, parameter, private :: debug = .false.

  ! estimate jacobian by differences or derivations ?
  logical, parameter, private :: analytic = .true.

  ! weights?
  logical, private :: weights

  integer, private :: ndat
  real(dbl), dimension(:), allocatable, private :: xdata, ydata, xsig, ysig, &
       res, rho, psi, dpsi
  real(dbl), private :: rsig

  private :: absline, absfun, nscale, funder, fundif, difjac, loglikely, residuals, &
       newline, logline


contains

  subroutine rline(x,y,a,b,da,db,dx,dy,sigma,reliable,verbose)

    use medians

    real(dbl), dimension(:), intent(in) :: x,y
    real(dbl), intent(out) :: a,b
    real(dbl), dimension(:), intent(in), optional :: dx,dy
    real(dbl), intent(out), optional :: da,db
    real(dbl), intent(out), optional :: sigma
    logical, intent(out), optional :: reliable
    logical, intent(in), optional :: verbose

    real(dbl) :: s,d,sig,xcen,ycen,d1,d2

    logical :: abreli, sreli, reli
    integer :: n

    weights = present(dx) .or. present(dy)

    if( present(verbose) ) then
       xverbose = verbose
    else
       xverbose = .false.
    end if

    if( weights .and. .not. present(dy) ) stop 'rline(): dy must be present.'

    if( size(x) /= size(y) ) stop 'rline(): Size of x,y data sets are different.'

    if( present(dx) ) then
       if( size(dx) /= size(x) ) stop 'rline(): Size of x and dx are different.'
       if( .not. all(dx >= 0) ) stop 'rline(): all dx >= 0, it is required.'
    end if

    if( present(dy) ) then
       if( size(dy) /= size(x) ) stop 'rline(): Size of y and dy are different.'
       if( .not. all(dy > 0) ) stop 'rline(): all dy > 0, it is required.'
    end if

    n = size(x)

    if( n < 2 ) then
       a = 0
       b = 0
       if( present(da) ) da = 0
       if( present(db) ) db = 0
       if( present(sigma) ) sigma = 1
       if( present(reliable) ) reliable = .false.
       return

    else if( n == 2 ) then
       d = x(2) - x(1)
       if( abs(d) > 0 ) then
          b = (y(2) - y(1)) / (x(2) - x(1))
          a = (y(1) + y(1)) / 2 + b*(x(1) + x(2)) / 2
       else
          a = 0
          b = 0
       end if
       if( present(da) ) da = 0
       if( present(db) ) db = 0
       if( present(sigma) ) sigma = 1
       if( present(reliable) ) reliable = .false.
       return

    end if

    allocate(xdata(n),ydata(n),xsig(n),ysig(n),res(n),rho(n),psi(n),dpsi(n))
    ndat = n
    xdata = x
    ydata = y

    if( present(dx) ) then
       xsig = dx
    else
       xsig = 0.0_dbl
    end if

    if( present(dy) ) then
       ysig = dy
    else
       ysig = 1.0_dbl
    end if

    ! compute data centroids
    xcen = median(x)
    ycen = median(y)
    if( xverbose ) write(*,*) "Centroids:",real(xcen),real(ycen)

    ! inital estimate by absolute deviations
    call absline(a,b,s,abreli)
    rsig = s
    call nscale(a,b,s,sreli)
    if( xverbose ) write(*,'(a,1p,2g15.5,0pf8.2,1x,2l1)') &
         'Robust line by abs.dev., sig by entropy:',a,b,s,abreli,sreli

    if( sreli ) rsig = s
    reli = abreli .and. sreli
    if( reli ) then

       call logline(a,b,abreli)
       call nscale(a,b,s,sreli)
       if( xverbose ) write(*,'(a,1p,2g15.5,0pf8.2,1x,2l1)') &
            'Robust line by likelihood, sig by entropy:',a,b,s,abreli,sreli

       if( sreli ) rsig = s
       reli = abreli .and. sreli
       if( reli ) then

          ! hight precise solution and statistical errors
          call newline(a,b,d1,d2,sig,reli)

       end if

    end if

    if( .not. reli ) then
       d2 = rsig/sqrt(real(n))
       if( b > 0 ) then
          d1 = d2 / b
       else
          d1 = huge(d1)
       end if
    end if

    if( present(da) ) da = d1
    if( present(db) ) db = d2
    if( present(reliable) ) reliable = reli
    if( present(sigma) ) sigma = sig

    deallocate(xdata,ydata,xsig,ysig,res,rho,psi,dpsi)

  end subroutine rline

  subroutine absline(a,b,s,reli)

    use NelderMead
    use medians

    real(dbl), intent(out) :: a,b,s
    real(dbl), dimension(2) :: t, dt
    logical, intent(out) :: reli
    real(dbl) :: ams, mad
    integer :: ifault

    t = 0
    dt = 1
    call nelmin1(absfun,t,dt,ams,ifault)
    reli = ifault == 0

    if( xverbose .and. ifault > 0 ) write(*,*) "No konvergence in rline: absline()."

    call residuals(t(1),t(2),res)

    if( ndat < 50 ) then
       mad = median(abs(res))
    else
       mad = qmedian(abs(res))
    end if

    a = t(1)
    b = t(2)
    s = mad / 0.6745

  end subroutine absline

  function absfun(t) result(s)

    real(dbl), dimension(:), intent(in) :: t
    real(dbl) :: s

    call residuals(t(1),t(2),res)
    s = sum(abs(res))

  end function absfun


  subroutine logline(a,b,reli)

    use NelderMead

    real(dbl), intent(in out) :: a,b
    real(dbl), dimension(2) :: t, dt
    logical, intent(out) :: reli
    real(dbl) :: ams
    integer :: ifault

    t = (/ a, b /)
    dt = rsig
    call nelmin1(loglikely,t,dt,ams,ifault)
    reli = ifault == 0

    if( xverbose .and. ifault > 0 ) write(*,*) "No konvergence in rline / logline."

    a = t(1)
    b = t(2)

  end subroutine logline


  subroutine newline(a,b,da,db,sig,reli)

    use minpacks

    real(dbl), intent(in out) :: a,b
    real(dbl), intent(out) :: da,db,sig
    real(dbl), dimension(2) :: t
    logical, intent(out) :: reli
    real(dbl), dimension(2,2) :: cov
    real(dbl) :: df, f2
    integer :: nprint,info, i

    if( debug ) then
       nprint = 1
    else
       nprint = 0
    end if

    t = (/ a, b /)

    ! Try to find the best solution. Our initial estimate is very neary
    ! to the right one.
    if( analytic ) then
       call lmder3(funder,t,cov,nprint=nprint,info=info)
    else
       call lmdif3(fundif,t,cov,nprint=nprint,info=info)
    end if

    reli = info /= 5
    if( reli ) then
       ! the minimum localised successfully, update estimates
       a = t(1)
       b = t(2)

       ! estimate uncerainities
       call residuals(a,b,res)
       call tukeys(res/rsig,psi)
       call dtukeys(res/rsig,dpsi)
!       dpsi = 1
!       psi = res
       df = sum(dpsi)
       f2 = sum(psi**2)

       if( df > 0 .and. f2 > 0 ) then

          cov = cov / rsig**2
          sig = rsig * sqrt(f2/df*ndat/(ndat-2))
          if( cov(1,1) > 0 ) da = sig * sqrt(cov(1,1))
          if( cov(2,2) > 0 ) db = sig * sqrt(cov(2,2))

       else
          reli = .false.
       end if

       if( xverbose ) then
          write(*,*) 'cov:',real(cov(1,:))
          write(*,*) 'cov:',real(cov(2,:))
       end if

    end if

  end subroutine newline


  subroutine residuals(a,b,r)

    real(dbl), intent(in) :: a,b
    real(dbl), dimension(:), intent(out) :: r

    if( weights ) then
       r = (ydata - (a + b*xdata)) / sqrt(ysig**2 + b**2*xsig**2)
    else
       r = ydata - (a + b*xdata)
    end if

  end subroutine residuals

  function loglikely(t)

    use rfun

    real(dbl), dimension(:), intent(in) :: t
    real(dbl) :: loglikely

    call residuals(t(1),t(2),res)
    call itukeys(res/rsig,rho)

    loglikely = sum(rho)

  end function loglikely

  subroutine funder(m,np,p,fvec,fjac,ldfjac,iflag)

    use rfun

    integer, intent(in) :: m,np,ldfjac
    integer, intent(inout) :: iflag
    real(dbl), dimension(np), intent(in) :: p
    real(dbl), dimension(m), intent(out) :: fvec
    real(dbl), dimension(ldfjac,np), intent(out) :: fjac
    real(dbl), dimension(:), allocatable :: r,ds,dra,drb
    real(dbl), dimension(2) :: fv
    real(dbl), dimension(2,2) :: dfjac
    real(dbl) :: a,b,s
    integer :: n

    if( iflag == 0 ) then

       write(*,'(4g15.5)') p,fvec

       if( debug ) then
          write(*,*) 'djac:',fjac(1,:)
          write(*,*) 'djac:',fjac(2,:)

          call difjac(p(1),p(2),dfjac)
          write(*,*) 'djac:',dfjac(1,:)
          write(*,*) 'djac:',dfjac(2,:)
       end if

       return
    end if

    n = size(xdata)
    allocate(r(n),ds(n))

    a = p(1)
    b = p(2)
    s = rsig

    if( weights ) then

       ds = sqrt(ysig**2 + b**2*xsig**2)
       r = (ydata - (a + b*xdata)) / ds
       call tukeys(r/s,psi)
!       psi = r/s
       fv(1) = sum(psi/ds)
       fv(2) = sum(psi*(xdata/ds + r*b*xsig**2/ds**2))

    else

       r = ydata - a - b*xdata
       call tukeys(r/s,psi)
!       psi = r / s
       fv(1) = sum(psi)
       fv(2) = sum(psi*xdata)

    end if

    if( iflag == 1 ) then

       fvec = - fv / s
!       write(*,*) fvec

    else if( iflag == 2 ) then

       allocate(dra(n),drb(n))
       call dtukeys(r/s,dpsi)
!       dpsi = 1

       if( weights ) then

          dra = - 1 / ds
          drb = - (xdata*ds + r*b*xsig**2) / ds**2

          fjac(1,1) = sum(dpsi/ds**2)
          fjac(1,2) = sum(dpsi*dra*drb) - s*sum(psi*b*xsig**2/ds**3)
          fjac(2,1) = fjac(1,2)
          fjac(2,2) = sum(dpsi*drb**2) - &
               s*sum(psi*(r*(1-3*b**2**xsig**2/ds**2)-2*xdata/ds)*xsig**2/ds**2)

       else

          fjac(1,1) = sum(dpsi)
          fjac(1,2) = sum(dpsi*xdata)
          fjac(2,1) = fjac(1,2)
          fjac(2,2) = sum(dpsi*xdata**2)

       end if

       fjac = fjac / s**2

       deallocate(dra,drb)

    end if

    deallocate(r,ds)

  end subroutine funder

  subroutine difjac(t,s,jac)

    real(dbl), intent(in) :: t,s
    real(dbl), dimension(:,:), intent(out) :: jac
    real(dbl), parameter :: d = 1e-4
    real(dbl), dimension(2) :: fv1,fv2
    integer :: iflag

    iflag = 1

    call fundif(2,2,(/t+d,s/),fv1,iflag)
    call fundif(2,2,(/t-d,s/),fv2,iflag)
    jac(1,:) = (fv1 - fv2)/(2*d)

    call fundif(2,2,(/t,s+d/),fv1,iflag)
    call fundif(2,2,(/t,s-d/),fv2,iflag)
    jac(2,:) = (fv1 - fv2)/(2*d)

  end subroutine difjac

  subroutine fundif(m,np,p,fvec,iflag)

    use rfun

    integer, intent(in) :: m,np
    integer, intent(in out) :: iflag
    real(dbl), dimension(np), intent(in) :: p
    real(dbl), dimension(m), intent(out) :: fvec
    real(dbl), dimension(:), allocatable :: r,f,ds
    real(dbl), dimension(2,2) :: jac
    real(dbl), dimension(2) :: fv
    integer :: n
    real(dbl) :: a,b,s

    if( iflag == 0 ) then

       write(*,'(6g13.5)') p,fvec

       if( debug ) then
          n = 2
          call funder(2,2,p,fv,jac,2,n)
          write(*,*) ' jac:',jac(1,:)
          write(*,*) ' jac:',jac(2,:)
       end if

       return
    end if

    n = size(xdata)
    allocate(r(n),f(n),ds(n))

    a = p(1)
    b = p(2)
    s = rsig

    if( weights ) then

       ds = sqrt(ysig**2 + b**2*xsig**2)
       r = (ydata - (a + b*xdata)) / ds
       call tukeys(r/s,psi)
!       psi = r/s

       fvec(1) = sum(psi/ds)
       fvec(2) = sum(psi*(xdata*ds + r*b*xsig**2)/ds**2)

    else

       r = ydata - (a + b*xdata)
       call tukeys(r/s,psi)
!       psi = r / s
       fvec(1) = sum(psi)
       fvec(2) = sum(psi*xdata)

    end if

    fvec = - fvec / s

    deallocate(r,f,ds)

  end subroutine fundif

  subroutine nscale(a,b,s,reli)

    use entropyscale

    real(dbl), intent(in) :: a,b
    real(dbl), intent(in out) :: s
    logical, intent(out) :: reli

    call residuals(a,b,res)
    call escale(res,s,reli)

  end subroutine nscale


end module robustline
