!
!  Robust Flat     Average of normalised flat-fields
!  Copyright (C) 2017  Filip Hroch, Masaryk University, Brno, 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 robustflat

  use iso_fortran_env

  implicit none

  logical, private :: verb = .false.

  integer, parameter, private :: dbl = selected_real_kind(15)

  real(dbl), dimension(:), allocatable, private :: flat, dflat, cf, dcf, rat
  real(dbl), private :: reg, sigma, df, rmin, rmax

  private :: residuals, fmean, zfun, fscale

contains

  subroutine rflat(flats,dflats,factors,dfactors,fin,dfin,reliable,verbose)

    use medians

    real, dimension(:), intent(in) :: flats,dflats,factors,dfactors
    real, intent(in out) :: fin, dfin
    logical, intent(out) :: reliable
    logical, intent(in), optional :: verbose

    real(dbl) :: f,sigold,regold,sig,fold
    integer :: n,i
    logical :: freli,sreli,convergent

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

    n = size(flats)

    if( n == 0 ) then
       fin = 1
       dfin = 0
       reliable = .false.
       return
    end if

    if( n <= 2 ) then
       fin = sum(flats / factors) / n
       dfin = sum(abs(flats / factors - fin)) / n
       reliable = .false.
       return
    end if

    ! Now, we works with appropriate amount of data: n > 2
    allocate(flat(n),dflat(n),cf(n),dcf(n),rat(n))
    flat = flats
    dflat = dflats
    cf = factors
    dcf = dfactors
    f = fin
    df = dfin

    rat = flat / cf
    rmin = minval(rat)
    rmax = maxval(rat)

    ! initial estimate of parameters
    call finit(f,sig)
    if( verb ) &
         write(error_unit,'(a,f0.5,1x,f0.3)') 'rflat initial f,sig: ',f,sig
    fold = f

    ! avoid other computations on identical data, or
    ! low amount of data for robust estimation
    if( abs(sig) < 2*epsilon(sig) ) then
       fin = real(f)
       dfin = 0
       reliable = .true.
       goto 666
    end if

    sigold = -1
    do i = 0,666
       reg = i
       call fmean(f,sig,freli)
       call fscale(f,sig,sreli)
       if( verb ) write(error_unit,'(a,f12.7,2(1x,f0.3),2x,2l)') &
            '      zmean,sig,reg:',f,sig,reg,freli,sreli
       reliable = freli .and. sreli
       if( sig < 1 ) exit
       sigold = sig
       regold = reg
    end do
    ! solve reg. parameter by linear interpolation to the condition sig=1
    if( sigold > 0 .and. abs(sig - sigold) > epsilon(sig) ) then
       reg = regold + (1 - sigold)*(reg - regold)/(sig - sigold)
       call fmean(f,sig,reliable)
    end if
    if( verb ) write(error_unit,'(a,f12.7,2(1x,f0.3),2x,l)') &
         'final zmean,sig,reg:',f,sig,reg,reliable

    if( reliable ) then
       ! Newton's update
       call fnewton(f,df,sig,convergent)
       if( verb ) write(error_unit,'(a,2f13.7,2x,l)') 'Newton:',f,df,convergent
    else
       fin = real(fold)
    end if

    if( convergent .and. reliable ) then
       fin = real(f)
       dfin = real(df)
    end if

666 continue
    deallocate(flat,dflat,cf,dcf,rat)

  end subroutine rflat

  subroutine residuals(f,res,sig)

    real(dbl), intent(in) :: f
    real(dbl), dimension(:), intent(out) :: res,sig

    res = flat - cf*f
    sig = sqrt(dflat**2 + dcf**2 + flat*reg**2)

  end subroutine residuals


  subroutine fmean(f,sig,reli)

    use fmm
    use medians

    ! This is a robust estimation of the flat.

    real(dbl), intent(in out) :: f
    real(dbl), intent(in) :: sig
    logical, intent(out) :: reli
    real(dbl) :: htol,h,d,gmin,gmax

    sigma = sig
    h = 3.3 * median(abs(rat - f)) / 0.6745
    gmin = max(f - h, rmin)
    gmax = min(f + h, rmax)
    htol = max(1e-6,1e-5*h)
!    if( verb ) write(error_unit,*) f,gmin,gmax,h,htol,log((gmax-gmin)/htol)

    f = zeroin(gmin,gmax,zfun,htol)

    d = 2*htol
    reli = gmin + d < f .and. f < gmax - d

    ! the realiable result is inside interval rmin .. rmax
    ! and not too close to the interval endpoints

  end subroutine fmean


  function zfun(f)

    use rfun

    real(dbl), intent(in) :: f
    real(dbl) :: zfun
    real(dbl), dimension(:), allocatable :: psi,res,sig
    integer :: n

    n = size(flat)
    allocate(psi(n),res(n),sig(n))
    call residuals(f,res,sig)
    call tukeys(res/(sigma*sig),psi)

    zfun = -sum(psi*cf/sig)

    deallocate(psi,res,sig)

  end function zfun

  subroutine fscale(f,s,reli)

    use entropyscale

    real(dbl), intent(in) :: f
    real(dbl), intent(in out) :: s
    logical, intent(out) :: reli
    real(dbl), dimension(:), allocatable :: res,sig
    integer :: n

    n = size(flat)
    allocate(res(n),sig(n))
    call residuals(f,res,sig)
    call escale(res/sig,s,reli)
    deallocate(res,sig)

  end subroutine fscale

  subroutine finit(f,s)

    use medians

    real(dbl), intent(in out) :: f
    real(dbl), intent(out) :: s

    real(dbl), dimension(:), allocatable :: res,sig
    integer :: n

    n = size(flat)
    f = median(rat)
    allocate(res(n),sig(n))
    call residuals(f,res,sig)
    s = median(abs(res/sig)) / 0.6745
    deallocate(res,sig)

  end subroutine finit

  subroutine fnewton(f,df,s,convergent)

    use rfun

    real(dbl), intent(in out) :: f
    real(dbl), intent(out) :: df
    real(dbl), intent(in) :: s
    logical, intent(out) :: convergent

    real(dbl), dimension(:), allocatable :: psi,dpsi,res,sig,rs
    real(dbl) :: fun, dfun, fun2, d, xtol
    integer :: iter,n

    convergent = .false.

    if( .not. (s > 0) ) return

    n = size(flat)
    allocate(psi(n),dpsi(n),res(n),sig(n),rs(n))

    ! tolerance limit depends on number of elements in sums
    ! because every element can include its rounding error...
    ! ... proportional to absolute value ...
    xtol = 10*n*abs(f)*epsilon(f)

    do iter = 1, precision(f)

       call residuals(f,res,sig)

       rs = res / (s * sig)
       call tukeys(rs,psi)
       call dtukeys(rs,dpsi)

       fun = -sum(psi*cf/sig)
       dfun = sum(dpsi*(cf/sig)**2)

       if( abs(dfun) < epsilon(dfun) ) exit

       ! corrector for mean
       d = s * fun / dfun

       ! update location
       f = f - d

       if( verb ) write(error_unit,'(a,i2,1pg20.10,1p3g12.3)') &
            "mean, incr., f, f': ",iter,f,d,fun,dfun

       convergent = abs(d) < xtol

       ! exit of iterations: the absolute errot must be at least |d| < tol
       if( convergent ) exit

    end do

    fun2 = sum(psi**2)  ! mean of psi**2
    dfun = dfun * s**2
    if( convergent .and. dfun > 0 .and. fun2 > 0 ) then
       df = s*sqrt(fun2/dfun**2*n/(n-1)) ! df is an element of Jacobian
    else
       df = s/sqrt(real(n))
    end if

    deallocate(res,sig,psi,dpsi,rs)

  end subroutine fnewton

end module robustflat
