!
! robfit - robust fitting
!
! Copyright © 2011-8 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 robfit

  use iso_fortran_env

  implicit none

  logical, private :: debug = .false.
  logical, private :: analytic = .true.
  logical, private :: rnewton = .true.
  integer, parameter, private :: dbl = selected_real_kind(15)
  real(dbl), parameter, private :: rad = 57.295779513082322865_dbl
  real(dbl), dimension(:), allocatable, private :: uref,vref
  real(dbl), dimension(:), pointer, private :: xstar,ystar
  real(dbl), private :: refl, xcen, ycen, scale
  integer, private :: ndat

  private :: dres, minfun, fundif, funder, difjac, difvec, hessian_min

contains

  subroutine robmin(type,a,d,xc,yc,rf,xx,yy,mad,acen,dacen,dcen,ddcen, &
       sc,dsc,pa,dpa,stdsig,verbose,reliable)

    ! find proper transformation by robust method

    use oakleaf
    use astrotrafo
    use minpacks
    use neldermead

    ! parameters
    character(len=*),intent(in) :: type
    real(dbl),dimension(:), intent(in) :: a, d
    real(dbl),dimension(:), target, intent(in) :: xx,yy
    real(dbl),intent(in) :: mad,xc,yc,rf
    real(dbl),intent(inout) :: acen,dcen,sc,pa
    real(dbl),intent(out) :: dacen,ddcen,dsc,dpa,stdsig
    logical, intent(in) :: verbose
    logical, intent(out) :: reliable

    real(dbl), parameter :: eps = 10*epsilon(1.0_dbl)
    integer, parameter :: npar = 4

    real(dbl),dimension(npar) :: p,dp,p0
    real(dbl),dimension(npar,npar) :: fjac,cov
    real(dbl),dimension(:), allocatable :: du,dv,res
    integer :: iter,i,j,info,nprint,ifault,nra,ndec
    real(dbl) :: s,smin
    type(AstroTrafoProj) :: t,ti
    character(len=20) :: fmt
    character :: areli
    logical :: sreli, treli

    debug = verbose
    reliable = .false.

    if(debug) write(error_unit,*) '=== Robust fitting ==='

    stdsig = huge(stdsig)
    refl = rf
    xcen = xc
    ycen = yc
    ndat = size(a)
    xstar => xx
    ystar => yy

    if( .not. (ndat > npar) ) stop 'Robust fitting requires five stars > 5.'
    if( mad < epsilon(mad) ) stop 'Robust fit needs mad > 0.'

    allocate(uref(ndat),vref(ndat),du(ndat),dv(ndat),res(2*ndat))

    call trafo_init(t,type,acen,dcen)
    call proj(t,a,d,uref,vref)

    ! setup
    scale = mad / 0.6745
    p = [ 0.0_dbl, 0.0_dbl, 1.0_dbl / sc, pa / rad ]
    dp = [ 0.1*p(3), 0.1*p(3), 0.001*p(3), 0.1 / rad ]

    do iter = 1, precision(p)

       p0 = p

       if( rnewton ) then
          ! This piece of code has origin in the gold good times, when
          ! no divengence could had occured and the robust function
          ! has been on base Huber's minimax(). Tukey's functions
          ! needs more carefull approach and some plan B as fail-back.
          nprint = 0 ! or 1
          if( analytic ) then
             call lmder2(funder,p,eps,nprint,info)
          else
             call lmdif2(fundif,p,eps,nprint,info)
          end if
          if( info == 0 ) &
               write(error_unit,*) 'Error in robfit: Improper input parameters.'

          treli = info == 2 .or. info == 3
          ifault = info

       else
          treli = .false.
       end if

       if( .not. treli ) then

          p = p0
          ! Plan B
          ! Update estimate of parameters by method without derivation
          ! Because there are no methods reliable for finding of roots
          ! in multidimensional case, we replaced root finding by minimizing.
          call nelmin1(minfun,p,dp,smin,ifault)
          treli = ifault == 0
          if( .not. treli ) p = p0

       end if

       areli = ' '
       if( treli .and. iter > 1 ) then

          ! update estimate of dispersion by entropy,
          ! applied only on the already iterated center of projection
          call dres(p,res(1:ndat),res(ndat+1:2*ndat))
          call iscale(res,s,sreli)

          ! accept only reliable estimates
          if( sreli ) then
             scale = s
             areli = 'T'
          else
             areli = 'F'
          end if
       end if

       if( debug ) &
            write(error_unit,'(2i2,es10.2,2f11.5,2g12.3,f10.1,f7.1,1x,l1,a1)') &
            iter,ifault,scale,acen,dcen,p(1:2)/p(3),1/p(3),rad*p(4),treli,areli

       ! finish when parameters are appropriate small
       ! and iterations couldn't made further progress
       ! the precision is set to be better than 1/1000 of a pixel
       ! and reached usually within two iterations
       if( iter > 1 .and. all(abs(p(1:2)/p(3)) < 0.001) .and. treli ) then
          reliable = .true.
          exit
       end if

       ! update center of projection and projected coordinates
       call trafo_init(ti,type,acen=acen,dcen=dcen, &
            xcen=xcen,ycen=ycen,scale=p(3),rot=rad*p(4),refl=refl)
       call invproj(ti,p(1),p(2),acen,dcen)
       call trafo_init(t,type,acen,dcen)
       call proj(t,a,d,uref,vref)
       p(1:2) = 0.0_dbl

    end do

    ! statistical errors of parameters
    ! Hessian is recomputed analyticaly in view of the fact that
    ! we needs errors of acen, dcen (but p(1:2) are X0,Y0) and
    ! Jacobian by lmdif2 is scaled by errors (multiplication of s**2
    ! is satisfactory only for p(3:4).
    call hessian_min(p,fjac)
    call qrinv(fjac,cov)

    if( debug ) then
       write(fmt,'(a,i0,a)') '(a,',npar,'g15.5)'
       write(error_unit,'(a)') '# Hessian at minimum:'
       write(error_unit,fmt) ('#',fjac(i,:),i=1,npar)
       write(error_unit,'(a)') '# Covariance matrix (no regularisation):'
       write(error_unit,fmt) ('#',cov(i,:),i=1,npar)
       write(error_unit,'(a)') '# Correlation matrix:'
       do i = 1,npar
          write(error_unit,'(a)',advance="no")  '#'
          do j = 1,npar
             if( cov(i,i) > 0 .and. cov(j,j) > 0 ) then
                write(error_unit,'(f9.3)',advance="no") &
                     cov(i,j)/sqrt(cov(i,i)*cov(j,j))
             else
                ! Hessian is not regularised and Tukey's function has been used.
                write(error_unit,'(a9)',advance="no") ' *?* '
             end if
          end do
          write(error_unit,*)
       end do

    end if

    dp = -1
    do i = 1,npar
       if( scale > 0 .and. cov(i,i) > 0 ) &
            dp(i) = scale * sqrt(cov(i,i)/(ndat - npar))
    end do

    sc = 1.0_dbl/p(3)
    pa = rad*p(4)

    ! deviations of output parameters
    dacen = dp(1) / cos(dcen / rad) ! cos() sqeezes metric (distance) on sphere
    ddcen = dp(2)
    dsc = sc*(dp(3)/p(3))
    dpa = rad*dp(4)

    ! stdsig
    call dres(p,du,dv)
    stdsig = sum(tukey(du/scale)**2 + tukey(dv/scale)**2) / &
             sum(dtukey(du/scale) + dtukey(dv/scale))**2 * ndat / 2
    stdsig = scale * sqrt(stdsig)

    if( debug ) then
       call dres(p,du,dv)
       nra = count(du > 0)
       ndec = count(dv > 0)

       write(error_unit,'(a,4es13.3)') '# solution:  ',p(1:npar)
       write(error_unit,'(a,4es13.2)') '# deviations:',dp(1:npar)
       write(error_unit,'(a,es10.2,a,g0.3,a)') '# stdsig: ',stdsig, &
            ' [deg]    ',3600*stdsig,' [arcsec]'
       write(error_unit,'(a,i7,"/",f0.1,2(i5,"+-",f0.1))') &
            '# sign test (total/expected, RA+, Dec+): ',&
            ndat,ndat/2.0,nra,sqrt(nra*0.25),ndec,sqrt(ndec*0.25)
       ! simple version of sign test, var = n*p*(1-p), where p = 0.5
    end if

    deallocate(uref,vref,du,dv,res)

  end subroutine robmin


  subroutine dres(p,du,dv,u,v)

    use astrotrafo

    real(dbl), dimension(:), intent(in) :: p
    real(dbl), dimension(:), intent(out) :: du,dv
    real(dbl), dimension(:), intent(out), optional :: u,v

    real(dbl) :: uu,vv,x,y
    integer :: i
    type(AstroTrafoProj) :: t

    call trafo_init(t,xcen=xcen,ycen=ycen, &
         ucen=p(1),vcen=p(2),scale=p(3),rot=rad*p(4),refl=refl)

    do i = 1, ndat
       call invaffine(t,xstar(i),ystar(i),uu,vv,x,y)
       du(i) = uref(i) - uu
       dv(i) = vref(i) - vv

       ! u,v - projected coordinates of refence stars, refreshed in main loop
       ! x,y - coordinates of stars on CCD
       ! xx,yy - rotated (and reflexed), but not scaled, x,y coordinates
       ! du,dv - residuals
       ! uu,vv - rotated, sometimes reflexed, and scaled x,y coordinates

       if( present(u) .and. present(v) ) then
          u(i) = x
          v(i) = y
       end if
    end do

  end subroutine dres


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

    use oakleaf

    integer, intent(in) :: m,np
    integer, intent(inout) :: iflag
    real(dbl), dimension(np), intent(in) :: p
    real(dbl), dimension(m), intent(out) :: fvec
    real(dbl), dimension(:), allocatable :: dx,dy,xx,yy
    real(dbl), dimension(4) :: dfvec
    real(dbl) :: sf,sc,sx,sy,fx,fy
    integer :: i

    if( iflag == 0 .and. debug ) &
         write(error_unit,*) '#robfit t:',real(p)

    sf = 0.0_dbl
    sc = 0.0_dbl
    sx = 0.0_dbl
    sy = 0.0_dbl

    allocate(dx(ndat),dy(ndat),xx(ndat),yy(ndat))
    call dres(p,dx,dy,xx,yy)
    do i = 1, ndat
       fx = tukey(dx(i)/scale)
       fy = tukey(dy(i)/scale)
!       fx = huber(dx(i)/scale)
!       fy = huber(dy(i)/scale)
!       fx = dx(i)/scale
!       fy = dy(i)/scale
       sx = sx + fx
       sy = sy + fy
       sc = sc + (fx*xx(i) + fy*yy(i))
       sf = sf + (fx*yy(i) - fy*xx(i))
    end do
    fvec = [ -sx,-sy,-sc,sf*p(3) ] /scale
    deallocate(dx,dy,xx,yy)

    if( debug .and. .false. ) then
       write(error_unit,'(a,4g15.5)') '    ',fvec * scale
       call difvec(p,dfvec)
       write(error_unit,'(a,4g15.5)') 'dif:',dfvec * scale
    end if

  end subroutine fundif


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

    use oakleaf

    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 :: fu,fv,dfu,dfv,du,dv,u,v,ru,rv
    real(dbl), dimension(m) :: fc
    real(dbl), dimension(4,4) :: dfjac
    real(dbl), dimension(4) :: dfvec
    real(dbl) :: r
    integer :: i,j

    if( iflag == 0 ) then
       write(error_unit,'(4f15.5)') p
       return
    end if

    allocate(fu(ndat),fv(ndat),du(ndat),dv(ndat),u(ndat),v(ndat), &
         ru(ndat),rv(ndat))
    call dres(p,du,dv,u,v)

    r = p(3)
    ru = du / scale
    rv = dv / scale
    fu = tukey(ru)
    fv = tukey(rv)
!    fu = huber(ru)
!    fv = huber(rv)
!    fu = ru
!    fv = rv

    fc(1) =-sum(fu)
    fc(2) =-sum(fv)
    fc(3) =-sum(fu*u + fv*v)
    fc(4) = sum(fu*v - fv*u) * r

    if( iflag == 1 ) then

       fvec = fc / scale

       if( debug .and. .false. ) then
          write(error_unit,'(a,4g15.5)') '    ',fvec * scale
          call difvec(p,dfvec)
          write(error_unit,'(a,4g15.5)') 'dif:',dfvec * scale
       end if

    else if( iflag == 2 ) then

       allocate(dfu(ndat),dfv(ndat))

       dfu = dtukey(ru)
       dfv = dtukey(rv)
!       dfu = dhuber(ru)
!       dfv = dhuber(rv)
!       dfu = 1
!       dfv = 1

       fjac = 0.0_dbl
       fjac(1,1) = sum(dfu)
       fjac(1,3) = sum(dfu*u)
       fjac(1,4) =-sum(dfu*v)*r
       fjac(2,2) = sum(dfv)
       fjac(2,3) = sum(dfv*v)
       fjac(2,4) = sum(dfv*u)*r
       fjac(3,3) = sum(dfu*u**2 + dfv*v**2)
       fjac(3,4) = sum((dfv - dfu)*u*v)*r
       fjac(4,4) = sum(dfu*v**2 + dfv*u**2)*r**2
       fjac = fjac / scale**2

       fjac(3,4) = fjac(3,4) + fc(4) / r / scale
       fjac(4,4) = fjac(4,4) - fc(3) * r / scale

       do i = 1,np
          do j = 1,i-1
             fjac(i,j) = fjac(j,i)
          end do
       end do

       if( debug .and. .false. ) then
          do i = 1,np
             write(error_unit,'(4g15.5)') fjac(i,:) * scale**2
          end do
          call difjac(p,dfjac)
          do i = 1,np
             write(error_unit,'(a,4g15.5)') 'dif: ',dfjac(i,:) * scale**2
          end do
       end if

       deallocate(dfu,dfv)

    end if

    deallocate(fu,fv,du,dv,u,v,ru,rv)

  end subroutine funder

  subroutine hessian_min(p,hess)

    real(dbl), dimension(:), intent(in) :: p
    real(dbl), dimension(:,:), intent(out) :: hess
    real(dbl), dimension(size(p),size(p)) :: fjac
    real(dbl), dimension(size(p)) :: fvec
    integer :: iflag,m

    iflag = 2
    m = size(p)
    call funder(m,m,p,fvec,fjac,m,iflag)
    hess = fjac*scale**2

  end subroutine hessian_min


  subroutine difjac(p,jac)

    ! numerical approximation of jacobian
    ! https://en.wikipedia.org/wiki/Finite_difference

    real(dbl), dimension(:), intent(in) :: p
    real(dbl), dimension(:,:), intent(out) :: jac
    real(dbl), dimension(size(p)) :: d,fv1,fv2
    integer :: i,n,iflag
    real(dbl) :: h

    n = size(p)
    h = sqrt(epsilon(h))
    d = 0.0_dbl

    jac = 0
    do i = 1, n
       if( abs(p(i)) > epsilon(p) ) then
          h = sqrt(epsilon(h))*abs(p(i))
       else
          h = sqrt(epsilon(h))
       end if
       d(i) = h
       call fundif(n,n,p+d,fv1,iflag)
       call fundif(n,n,p-d,fv2,iflag)
       jac(i,i:n) = (fv1(i:n) - fv2(i:n))/(2*h)
       d(i) = 0.0_dbl
    end do


  end subroutine difjac

  subroutine difvec(p,vec)

    ! numerical approximation of jacobian
    ! https://en.wikipedia.org/wiki/Finite_difference

    real(dbl), dimension(:), intent(in) :: p
    real(dbl), dimension(:), intent(out) :: vec
    real(dbl), dimension(size(p)) :: d
    integer :: i,n
    real(dbl) :: h,f1,f2

    n = size(p)
    h = sqrt(epsilon(h))
    d = 0.0_dbl

    do i = 1,n
       d(i) = h
       f1 = minfun(p+d)
       f2 = minfun(p-d)
       vec(i) = (f1 - f2)/(2*h)
       d(i) = 0.0_dbl
    end do

  end subroutine difvec


  function minfun(p)

    use oakleaf

    real(dbl) :: minfun
    real(dbl), dimension(:), intent(in) :: p
    real(dbl), dimension(:), allocatable :: du,dv

    allocate(du(ndat),dv(ndat))
    call dres(p,du,dv)
    minfun = sum(itukey(du/scale)) / ndat + sum(itukey(dv/scale)) / ndat
!    minfun = (sum(ihuber(du/scale)) + sum(ihuber(dv/scale))) / (2*ndat)
!    minfun = (sum((du/scale)**2 / 2) + sum((dv/scale)**2 / 2)) / (2*ndat)

  end function minfun

end module robfit
