;
; maxtrace.pro - Finds and fits the peak along a trace
;
; Copyright (C) 2005, 2006 Patricio Rojo
;
; This program 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 2
; of the License, or (at your option) any later version.
;
; This program 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, write to the Free Software
; Foundation, Inc., 51 Franklin Street, Fifth Floor, 
;                   Boston, MA 02110-1301, USA.
;

;+
; NAME:
;	MAXTRACE
;
; PURPOSE:
;	This function fits the trace of some predominant features
;	along columns in a two dimensional array.
;
; CATEGORY:
;	Pato's Array Data Reduction Potpourri.
;
; CALLING SEQUENCE:
;
;	trace = MAXTRACE(Data, Y, Refy, Refx)
;
; INPUTS:
;       Data:   2 dimensional array in where to find the trace along
;               the columns.
;       Y:      Array containing the value of the y axis.
;       Refy:   Approximate value of Y where to start looking for the
;               trace.
;       Refx:   Approximate value of the X-coordinate (which is always
;               equivalent to index value) where to start looking for
;               the trace.
;
;
; KEYWORD PARAMETERS:
;       OSAMP:    Oversampling for peak finding. If using this option
;                 and a wavelet as a input you will want to use expy
;                 and be careful about converting back the output. If
;                 you want to use this switch make sure that the y
;                 scale is linear and equispaced.
;       EXPY:     Process through the natural logarithm of y
;                 coordinate.
;       FITNAME:  Name of the fit you want to use. Currently
;                 'true', 'GFCest', 'GVC', 'GVCnb',
;                 'GFCnb', and 'GFC'. Check defringeflat.pro
;                 documentation for a more detailed explanation of
;                 each of the methods. Default is 'GFC'.
;       NOVAL:    Value of parameters when there was no succesfull fit
;       IGNORE:   Ignore data beyond this columns
;       TRACELIM: Force trace to be within these values, if it is
;                 going beyond this value, it goes back to REFY value
;                 and start fitting in the opposite direction. Note
;                 that this keyword has not been tested thoroughly,
;                 please report any bug that you might find.
;       COIBORDERS: Stops looking for trace outside these borders, this
;                 array is in the COI format (cone of influence of a
;                 wavelet). If the trace's Y value goes above the
;                 corresponding value in this array, then computation
;                 stop.
;       BLINDCONTINUE: Continue the specified number of column blindly
;                 if a peak was not successfully found. It uses the
;                 last known trace center to find the next one. If
;                 none of this extra columns gives a successfull fit
;                 then it stop looking in that direction. A value of
;                 zero deactivates this alternative (Default).
;       STOPATSTART: Force stop at the start of this routine. Keyword
;                 for debugging purposes only.
;
; OUTPUTS:
;	This function returns the parameters of the fitted trace
;
; OPTIONAL OUTPUTS:
;       LOCALMINIMA: 2x#Y elements array containing the index of the
;                 local minima for each column
;       TRUEBACK: Array containing the background when using trueshape
;                 fitting. 
;       PEAKS:    Array containing the Y-value of the highest peaks in
;                 the trace along the columns.
;
; SIDE EFFECTS:
;
; RESTRICTIONS:
;
; PROCEDURE:
;	Algorithm description.
;
; EXAMPLE:
;      Find the trace of array 'arr' using default Gaussian fitting,
;   processing the background, given axis y is monospaced in
;   logarithmic space. Use position (cntx, xnty) as a reference to
;   start looking for the peak.
;
;       cntx = 220
;       cnty = 430
;	gaussprm = MAXTRACE(arr, y, cntx, cnty, /expy)
;
;
; MODIFICATION HISTORY:
; 	Written by:	Pato Rojo, Cornell.  Jul 2005 - Mar 2006
;			pato@astro.cornell.edu
;-

FUNCTION peakborders, y, data, cy
COMPILE_OPT idl2, hidden

n = N_ELEMENTS(y)
der = DBLARR(n)
sder = DBLARR(n)

der[2:n-3] = (data[3:n-2] - data[1:n-4]) / 2.0 + $
  (data[4:n-1] - data[0:n-5]) / 4.0

iyy = (WHERE(y GE cy))[0]

;if it is at the right side of the peak
IF der[iyy] LT 0 THEN BEGIN
    lind = WHERE(der[iyy:*] GE 0)
    IF N_ELEMENTS(lind) EQ 1 THEN u = n - 1 $
    ELSE BEGIN
        pos = iyy + lind
        WHILE(pos[1] EQ pos[0] + 2) DO               $   
          IF N_ELEMENTS(pos) GT 2 THEN pos = pos[1:*] $
          ELSE pos = [n-1, n-1]
        u = pos[0]
    ENDELSE
    lind = WHERE(der[0:iyy] GE 0)
    IF N_ELEMENTS(lind) EQ 1 THEN RETURN, [0, 1]  $ ;end of search
    ELSE BEGIN
        pos = ROTATE(lind, 2)
        WHILE(pos[1] EQ pos[0] - 2) DO                $
          IF N_ELEMENTS(pos) GT 2 THEN pos = pos[1:*] $
          ELSE RETURN, [0, 1]
        lind = WHERE(der[0:pos[0]-1] LT 0)
        IF N_ELEMENTS(lind) EQ 1 THEN l = 0 $
        ELSE BEGIN
            pos = ROTATE(lind, 2)
            WHILE(pos[1] EQ pos[0] - 2) DO                $
              IF N_ELEMENTS(pos) GT 2 THEN pos = pos[1:*] $
              ELSE pos = [0, 0]
            l = pos[0]
        ENDELSE
    ENDELSE
;if it is at the left
ENDIF ELSE BEGIN
    lind = WHERE(der[0:iyy] LT 0)
    IF N_ELEMENTS(lind) EQ 1 THEN l = 0 $
    ELSE BEGIN
        pos = ROTATE(lind, 2)
        WHILE(pos[1] EQ pos[0] - 2) DO                $
          IF N_ELEMENTS(pos) GT 2 THEN pos = pos[1:*] $
          ELSE pos = [0, 0]
        l = pos[0]
    ENDELSE
    lind = WHERE(der[iyy:*] lt 0)
    IF N_ELEMENTS(lind) EQ 1 THEN RETURN, [0, 1] $ ;end of search
    ELSE BEGIN
        pos = iyy + lind
        WHILE(pos[1] EQ pos[0] + 2) DO                $
          IF N_ELEMENTS(pos) GT 2 THEN pos = pos[1:*] $
          ELSE RETURN, [0, 1]
        lind = WHERE(der[pos[0]+1:*] GE 0)
        IF N_ELEMENTS(lind) EQ 1 THEN u = n - 1 $
        ELSE BEGIN
            pos = pos[0] + 1 + lind
            WHILE(pos[1] EQ pos[0] + 2) DO                $
              IF N_ELEMENTS(pos) GT 2 THEN pos = pos[1:*] $
              ELSE pos = [n-1, n-1]
            u = pos[0]
        ENDELSE
    ENDELSE
ENDELSE

RETURN, [l, u]

END

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

FUNCTION findtrace, refy, inix, finx, stepx, cyy, tracelim, ignore, $
                    y, data, yfine, n, trueback, osamp, dogaussfit, $
                    traceprm, coi, blindcontinue, carr

ocy = refy
l  = -1
contsteps = blindcontinue

;examining each column
FOR x = inix, finx, stepx DO BEGIN
    nextiter = 0
    IF(x GT ignore[1] OR x LT ignore[0]) THEN BREAK

    spdata = data[x, *]
    IF(osamp GT 1) THEN BEGIN
        sp     = SPL_INIT(y, spdata)
        spdata = SPL_INTERP(y, spdata, sp, yfine)
    ENDIF

    ;actualize old ones and find local minima
    ol = l
    cy = ocy
    l = PEAKBORDERS(yfine, spdata, cy)
    c = (LONG(l / FLOAT(osamp) + 0.5) < (N_ELEMENTS(y)-1)) > 0
    carr[*,x] = c

    IF(l[0] GE l[1]) THEN BEGIN 
        IF contsteps-- LE 0 THEN BREAK
        CONTINUE
    ENDIF

    tmpd = MAX(spdata[l[0]:l[1]], cntidx)
    cntidx += l[0]
    cy   = yfine[cntidx]

;stop looking for trace if new center is outside previous borders
    IF ol[0] NE -1 THEN $
      IF(cntidx GT ol[1] OR cntidx LT ol[0]) THEN BEGIN
        IF contsteps-- LE 0 THEN BREAK
        l = ol
        CONTINUE
    ENDIF

;stop looking if beyond coi borders
    IF cy GT coi[x] THEN BREAK

;if I'm backsearching then finish doing so when locating previous
;point
    IF N_ELEMENTS(tracelim) EQ 1 THEN BEGIN
        IF tracelim EQ 1 THEN BEGIN
            IF FLOAT(cyy[x]) EQ FLOAT(cy) THEN RETURN, x
        ENDIF ELSE IF tracelim NE 0 THEN MESSAGE, $
          "Value of tracelim is invalid, either 0, 1 or two elements"

;check whether I need to backtrace
    ENDIF ELSE IF (N_ELEMENTS(tracelim) EQ 2) THEN BEGIN
        IF ((cy GT tracelim[1]) OR (cy LT tracelim[0])) THEN BEGIN
            b = FINDTRACE(refy, x, inix, -stepx, cyy, 1, ignore, $
                          y, data, yfine, n, trueback, osamp,    $
                          dogaussfit, traceprm, coi, carr)
        ENDIF
    ENDIF ELSE $
      MESSAGE, "TRACELIM must be either 0, 1 or have two elements"

;update peaks
    cyy[x] = cy

    IF dogaussfit GT 0 THEN BEGIN
;prepare estimates
        wdt = yfine[l[1]] - yfine[l[0]]
        cnt = cy
        cidx = (WHERE(yfine GE cy))[0]
        gausest = [spdata[cidx], cnt, wdt*0.2, 0]
        s = 0
        func = ''

                                ;process with the fitting
        SWITCH dogaussfit OF
            1: BEGIN            ;Using fixed center no fitting
                lgaussprm = gausest
                BREAK
            END
            5: IF func EQ '' THEN BEGIN ;Use nb noncnt Gaussian
                a = [gausest[[0,1,2]]]
                func = 'fullgauss_nb'
                cnt = 0
            ENDIF
            4: IF func EQ '' THEN BEGIN ;Use nocnt Gaussian
                a = [gausest[[0,1,2,3]]]
                func = 'fullgauss'
                cnt = 0
            ENDIF
            3: IF func EQ '' THEN BEGIN ;Use nb fixed center Gaussian fitting
                a = [gausest[[0,2]]]
                func = 'evengauss_nb'
            ENDIF
            2: BEGIN            ;Using fixed center Gaussian fitting
                IF func EQ '' THEN BEGIN
                    a = [gausest[[0,2,3]]]
                    func = 'evengauss'
                ENDIF
                IF l[1]-l[0]+1 LE N_ELEMENTS(a) THEN BEGIN
                    IF contsteps-- LE 0 THEN BREAK
                    l = ol
                    nextiter = 1
                    BREAK
                ENDIF
                vx = yfine[l[0]:l[1]]

                gf = CURVEFIT(vx-cnt, spdata[l[0]:l[1]],               $
                              REPLICATE(1.0, l[1] - l[0] + 1), a,      $
                              STATUS=s, FUNCTI=func, ITER=iter)

;if background is less than zero, then force a fit without background
                IF (dogaussfit EQ 2) THEN IF (a[2] LT 0) THEN BEGIN
                    b  = a[0:1]
                    gf = CURVEFIT(vx-cnt, spdata[l[0]:l[1]],                $
                                  REPLICATE(1.0, l[1] - l[0] + 1), b,       $
                                  STAT=s, FUNC='evengauss_nb', ITER=iter)
                    a = [b, 0]
                ENDIF

;if background is less than zero, then force a fit without background
                IF (dogaussfit EQ 4) THEN IF (a[3] LT 0) THEN BEGIN
                    b  = a[0:2]
                    gf = CURVEFIT(vx-cnt, spdata[l[0]:l[1]],                $
                                  REPLICATE(1.0, l[1] - l[0] + 1), b,       $
                                  STAT=s, FUNC='fullgauss_nb', ITER=iter)
                    a = [b, 0]
                ENDIF

;If result doesn't make sense, then only use estimates
                IF (s EQ 1) OR (a[0] LE 0) OR $
                  (a[1] GE 50 * gausest[2]) THEN BEGIN
                    IF ~ N_ELEMENTS(lbase) THEN BEGIN
                        ocy = refy
                        nextiter = 1
                        BREAK
                    ENDIF

                    lgaussprm = [(traceprm[x-stepx, 0] +        $
                                  (gausest[0] - lbase)) / 2.0,  $
                                 cnt,                           $
                                 (traceprm[x-stepx, 2] +        $
                                  gausest[2]) / 2.0,            $
                                 lbase]
                ENDIF ELSE BEGIN
                    IF dogaussfit GT 3 THEN                    $
                      lgaussprm = [a[0], a[1], abs(a[2]),      $
                                   (dogaussfit EQ 4)?a[3]:0.0] $
                    ELSE                                       $
                      lgaussprm = [a[0], cnt, abs(a[1]),       $
                                   (dogaussfit EQ 2)?a[2]:0.0]
                ENDELSE

                BREAK
            END
            ELSE: MESSAGE, "Option not supported for Gaussian fitting" + $
              " 'dogaussfit'(" + STRING(dogaussfit) + ")"
        ENDSWITCH
        IF (s NE 0) THEN BEGIN
            IF contsteps-- LE 0 THEN BREAK
            l = ol
            CONTINUE
        ENDIF

;The following two conditions are checked above, hence they will only
;affect GVC
;stop looking for trace if new center is outside previous borders
        IF ol[0] NE -1 THEN BEGIN
            cntidx = (WHERE(yfine GT lgaussprm[1]))[0]
            IF(cntidx GT ol[1] OR cntidx LE ol[0]) THEN BEGIN
                IF contsteps-- LE 0 THEN BREAK
                l = ol
                CONTINUE
            ENDIF
        ENDIF
;stop looking if beyond coi borders
        IF cy GT coi[x] THEN BREAK

        IF nextiter THEN CONTINUE

        traceprm[x,*] = lgaussprm
        lbase = lgaussprm[3]
    ENDIF ELSE BEGIN
                                ;Using real shape with borders where
                                ;the neighboring peak starts
        bslope = (spdata[l[1]] - spdata[l[0]]) / (yfine[l[1]] - yfine[l[0]])
        binter = spdata[l[1]] - bslope * yfine[l[1]] 
        traceprm[x, c[0]:c[1]] = data[x, c[0]:c[1]]
        trueback[x, *] = [binter, bslope]
    ENDELSE
    ocy = cy
    contsteps = blindcontinue

ENDFOR

RETURN, x

END

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

PRO evengauss, x, A, F, pder

nx = x / a[1]
ex = exp(-0.5 * nx * nx)
f = a[0] * ex + a[2]

IF N_PARAMS() GE 4 THEN $
  pder = [[ex], [a[0] * ex * nx * nx / a[1]], $
          [replicate(1.0, n_elements(x))]]

END

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

PRO evengauss_nb, x, A, F, pder

nx = x / a[1]
ex = exp(-0.5 * nx * nx)
f = a[0] * ex

IF N_PARAMS() GE 4 THEN $
  pder = [[ex], [a[0] * ex * nx * nx / a[1]]]

END

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

PRO fullgauss, x, A, F, pder

nx = (x-a[1]) / a[2]
ex = exp(-0.5 * nx * nx)
f = a[0] * ex + a[3]

IF N_PARAMS() GE 4 THEN $
  pder = [[ex], [a[0] * ex * nx / a[2]], [a[0] * ex * nx * nx / a[2]], $
          [replicate(1.0, n_elements(x))]]

END

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

PRO fullgauss_nb, x, A, F, pder

nx = (x-a[1]) / a[2]
ex = exp(-0.5 * nx * nx)
f = a[0] * ex

IF N_PARAMS() GE 4 THEN $
  pder = [[ex], [a[0] * ex * nx / a[2]], [a[0] * ex * nx * nx / a[2]]]

END

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

FUNCTION maxtrace, data, y, refx, refy,                            $
                   OSAMP=osamp, FITNAME=fitname, EXPY=expy,        $
                   NOVAL=noval, IGNORE=ignore, TRUEBACK=trueback,  $
                   TRACELIM=tracelim, PEAKS=peaks,                 $
                   COIBORDERS=coiborders, STOPATSTART=stopatstart, $
                   BLINDCONTINUE=blindcontinue,                    $
                   LOCALMINIMA=localminima
;
;Finds trace of maximum point accross a continuous image
;
;Example:
;trace=maxtrace(data,y,refy,refx,/derivcenter,osamp=100,noval=0
;               ,/gaussianfit,/expy,ignore=[1,1000])
;

ON_ERROR,2
nx = (SIZE(data,/DIM))[0]
n  = (SIZE(data,/DIM))[1]
peaks = FLTARR(nx)

IF KEYWORD_SET(stopatstart) THEN stop

goback = -1

IF N_ELEMENTS(noval) EQ 0       THEN noval         = 0.0
IF N_ELEMENTS(tracelim) NE 2    THEN tracelim      = 0
if ~ N_ELEMENTS(ignore)         THEN ignore        = [0, nx]
IF ~ KEYWORD_SET(fitname)       THEN fitname       = 'GFC'
IF ~ KEYWORD_SET(blindcontinue) THEN blindcontinue = 0

IF N_ELEMENTS(y) NE n THEN $
  MESSAGE, "Y array needs to have the same size as the y-dimension" + $
  " of the data array" 
IF ~ KEYWORD_SET(coiborders) THEN coi      = DBLARR(nx) + y[n-1] $
ELSE BEGIN
    IF N_ELEMENTS(coiborders) EQ nx THEN coi = coiborders $
    ELSE MESSAGE, "coiborders array must have the same size as there " + $
      "is columns"
ENDELSE

IF(KEYWORD_SET(osamp)) THEN BEGIN
    IF KEYWORD_SET(expy) THEN yy = alog(y) $
    ELSE yy = y
    span    = yy[n-1] - yy[0]
    osamp = LONG(osamp)
    yfine   = yy[0] + span *             $
      DINDGEN(LONG(osamp * (n-1) + 1)) / DOUBLE(osamp * (n-1))
    IF KEYWORD_SET(expy) THEN yfine = EXP(yfine)
    n       = N_ELEMENTS(yfine)
ENDIF ELSE BEGIN
    yfine   = y
    osamp = 1
ENDELSE
lim = tracelim

dogaussfit = 0
SWITCH fitname OF
    'GVCnb': dogaussfit += 1
    'GVC': dogaussfit += 1
    'GFCnb': dogaussfit += 1
    'GFC':   dogaussfit += 1
    'GFCest': BEGIN
        dogaussfit += 1
        traceprm    = fltarr(nx, 4)
        nill        = replicate(noval, 4)
        BREAK
    END
    'true' : BEGIN
        dogaussfit = 0
        traceprm = dblarr(size(data,/dim))
        trueback = dblarr(nx, 2)
        nill     = replicate(noval, (size(data,/dim))[1])
        BREAK
    END
    ELSE: MESSAGE, "Unrecognized fit method '" + fitname + "'"
ENDSWITCH

ON_ERROR, 0
localminima = DBLARR(2, nx)

;Find the trace going forward and complete with nill what couldn't be
;fitted
x = FINDTRACE(refy, refx, nx-1, 1, peaks, lim, ignore, y, data,  $
              yfine, n, trueback, osamp, dogaussfit, traceprm, coi, $
              blindcontinue, localminima)
FOR x = x, nx-1, 1 DO traceprm[x, *] = nill

;Find the trace going backward and complete with nill what couldn't be
;fitted
x = FINDTRACE(refy, refx, 0, -1, peaks, lim, ignore, y, data,   $
              yfine, n, trueback, osamp, dogaussfit, traceprm, coi, $
              blindcontinue, localminima)
FOR x = x, 0, -1 DO traceprm[x, *] = nill

RETURN, traceprm

END