;---------------------------------------------------------------------------
; Document name: drift_rate.pro
; Created by:    Oliver Trachsel, January 2002
;
; Last Modified: 2002/01/30
;---------------------------------------------------------------------------
;
;+
; NAME:
;       Determine Drift Rate
;
; PURPOSE:
;	Calculate the drift rate by performing the following steps:
;		1) Cross-correlation on the currently displayed image
;		2) Spline fitting to add additional data points
;		3) Linear Regression of crosscorrelation maxima
;                  to calculate slope (drift_rate) and variance
;       After a first iteration the user can remove CC-maxima
;       with the mouse and the linear regression is then redone.	
;
; CALLING SEQUENCE
;	DriftRate, zReal, xAxis, yAxis, xMin, xMax, yMin, yMax, 
;                  addPts, /LOGFREQ
;
; INPUTS:
;       zReal: the 2D data array;
;	xAxis, yAxis: the two axes. 
;	xMin, xMax, yMin, yMax: the selected boundaries of
;		                the image.
;	addPts: Number of points to insert between existing ones
;               (by cubic spline interpolation).
;	
; OUTPUTS: 
;       None
;
; KEYWORDS:
;       LOGFREQ: If set, a logarithmic scale is used for frequency axis
;                and slope a is output instead of drift rate.
;       PSOUT:   If set, postscript files of the two (or one)
;                resulting plots will be created (named drift_plot.1.ps and
;                drift_plot_2.ps)
; 
; SIDE EFFECTS:
;
; RESTRICTIONS:
;
; MODIFICATION HISTORY:
;       02/01/30 Initial version
;                Oliver Trachsel <toliver@astro.phys.ethz.ch>
;		 Institute of Astronomy, ETH Zurich
;-

; utility procedure to swap contents of x and y
PRO SwapVars, x, y
    tmp = x
    x = y
    y = tmp
END

; utility procedure to plot results on active device
PRO Drift_Plot, ccImage, xAxis, spldXAxis, yAxis, yScale, maxima, $
                keepMaxima, const, slope, LOGFREQ=LOGFREQ
    xMax = N_ELEMENTS(xAxis)-1
    yMax = N_ELEMENTS(yAxis)-1
    maxValue = MAX(ABS(ccImage))
    IF KEYWORD_SET(LOGFREQ) THEN BEGIN
        ; Plot axes (to initialize data coordinates)
        YTicknames = REPLICATE(yAxis[0], 7) + INDGEN(7)* $
          ((yAxis[yMax]-yAxis[0])/6)
        YTicknames = STRTRIM(STRING(EXP(YTicknames), FORMAT='(I)'), 1)
        PLOT, [0], /NODATA, /YLOG, PSYM=-7, $
          XRANGE=[xAxis[0], xAxis[xMax]], XSTYLE=1, $
          YRANGE=[yAxis[0], yAxis[yMax]], YSTYLE=1, $
          YTICKNAME=YTicknames, YTICKS=N_ELEMENTS(YTicknames)-1
        ; Embed crosscorrelation image
        TV_Y_ADAPT,ccImage,xAxis,yAxis, MIN=-maxValue, MAX=maxValue
        ; Replot axis
        PLOT, [0], /NODATA, /YLOG, PSYM=-7, $
          XRANGE=[xAxis[0], xAxis[xMax]], XSTYLE=1, $
          YRANGE=[yAxis[0], yAxis[yMax]], YSTYLE=1, /NOERASE, $
          YTICKNAME=YTickNames, YTICKS=N_ELEMENTS(YTicknames)-1
        ; Plot CC maxima
        OPLOT, spldXAxis[maxima[keepMaxima]], yAxis[keepMaxima], $
          PSYM=-7, COLOR=0
    ENDIF ELSE BEGIN
        plot_image, ccImage, /NOSQUARE, MIN=-maxValue, MAX=maxValue, $
          ORIGIN=[xAxis[0], yAxis[0]], SCALE=[xAxis[1]-xAxis[0], yScale]
        OPLOT, spldXAxis[maxima[keepMaxima]], $
          yAxis[0]+DOUBLE(keepMaxima)*yScale, PSYM=-7, COLOR=0
    ENDELSE
    ; plot regression line
    OPLOT, const+slope*yAxis, yAxis, COLOR=0, THICK=3
END

PRO Drift_Rate, zReal, xAxis, yAxis, xMin, xMax, yMin, yMax, addPts, $
                LOGFREQ = LOGFREQ, PSOUT=PSOUT

; Perform cross-correlation
    
    ref_freq_index =  (yMax - yMin)/2

    correl1D, zReal[xMin:xMax, yMin:yMax], xAxis[xMin:xMax], $
      ref_freq_index, ccImage, ccXAxis
    
; Extract correlation maxima for each frequency

DO_SPLINE:
    maxima = LONARR(yMax-yMin+1) ; index of the maximum for each frequency
    IF addPts GT 0 THEN BEGIN
        cc_minimum = 0D
        cc_maximum = 0D
        
        ; Number of points before and after spline interpolation
        nofPts_bef = N_ELEMENTS(ccXAxis) ;
        nofPts_aft = nofPts_bef*(addPts+1) ; 
		
        newCcXAxis = DINDGEN(nofPts_aft)*(ccXAxis(nofPts_bef-1) - $
                                          ccXAxis(0))/nofPts_aft + ccXAxis(0)
		
        ; Spline Interpolation and maxima for each line

        FOR j=0L,(yMax-yMin) DO BEGIN
            splinedArr = SPLINE(ccXAxis, ccImage(*,j), newCcXAxis)
            tempMax = MAX(splinedArr, maxIndex, MIN=tempMin)
            maxima[j] = maxIndex
            IF tempMax GT cc_maximum THEN cc_maximum = tempMax
            IF tempMin LT cc_minimum THEN cc_minimum = tempMin
        ENDFOR
    ENDIF ELSE BEGIN
        FOR j=0L,(yMax-yMin) DO BEGIN
            tmp = MAX(ccImage(*, j), maxIndex)
            maxima(j) = maxIndex
        ENDFOR
        cc_maximum = MAX(ccImage, MIN=cc_minimum)
        newCcXAxis = ccXAxis
    ENDELSE

; Perform linear regression on array 'maxima'

    newYAxis = yAxis
    IF KEYWORD_SET(LOGFREQ) THEN newYAxis[ymin:ymax] = alog(yAxis[ymin:ymax])

    slope = REGRESS(TRANSPOSE(newYAxis[yMin:yMax]), newCcXAxis[maxima], $
                    REPLICATE(1.0, yMax-yMin+1), $
                    xfit, const, sigma, ftest, corr, rmul, chisq, $
                    /RELATIVE_WEIGHT)

; Display results & graphics    
    ny = ymax-ymin+1
    yScale = (yAxis[ymax]-yAxis[ymin])/ny

    IF N_ELEMENTS(winNr) EQ 0 THEN BEGIN
        oldWinNr = !D.WINDOW
        winNr = !D.WINDOW + 1
        WINDOW, winNr, TITLE = 'Linear regression (1st iteration)'
        oldPTitle = !P.TITLE
        oldXTitle = !X.TITLE
        oldYTitle = !Y.TITLE
        !P.TITLE = 'Crosscorrelation with maxima / Linear regression'
        !X.TITLE = 'time [s]'
        !Y.TITLE = KEYWORD_SET(LOGFREQ) ? 'log scale [MHz]' : 'frequency [MHz]'
    ENDIF ELSE BEGIN
        WSET, winNr
        WSHOW, winNr
    END

    Drift_Plot, ccImage, ccXAxis, newCcXAxis, newYAxis[yMin:yMax], yScale, $
      maxima, INDGEN(N_ELEMENTS(maxima)), const, slope[0], LOGFREQ=LOGFREQ

    ; do not reprint title if just spline interpolation has been
    ; recalculated
    IF N_ELEMENTS(isSecondRound) EQ 0 THEN BEGIN
        title = '  DRIFT RATE DETERMINATION'
        IF KEYWORD_SET(LOGFREQ) THEN title = title + ' (FREQUENCY IN LOG SCALE)'
        PRINT, '--------------------------------------------------------------------'
        PRINT, title
        PRINT, '--------------------------------------------------------------------'
    ENDIF
    PRINT
    PRINT, '  Spline interpolation:'
    PRINT, '    Points added between data points: ',STRN(addPts) 
    centerFreq = STRTRIM(string(yaxis[ymin+ref_freq_index], FORMAT='(D10.1)'),1)
    PRINT
    PRINT, '  Splined crosscorrelation:'
    PRINT, '	Center frequency: ', centerFreq, ' MHz'
    PRINT, '	Correlation minimum: ', STRN(cc_minimum)
    PRINT, '	Correlation maximum: ', STRN(cc_maximum)
    PRINT
    PRINT, '  Linear regression - first iteration:'
    PRINT, '	Sigma: ', STRN(sigma)
    PRINT, '	Corr: ', STRN(corr)
    PRINT, '	Chisq: ', STRN(chisq)
    PRINT
    IF KEYWORD_SET(LOGFREQ) THEN $
      PRINT, '  Slope: ',STRN(1./slope[0]),' +/- ',STRN(sigma/(slope[0]^2)), $
        ' [1/s]' $
    ELSE $
      PRINT, '  Drift rate: ',STRN(1./slope[0]),' +/- ', $
             STRN(sigma/(slope[0]^2)), ' [MHz/s] at ',centerFreq,' MHz'
    PRINT

; Give the user the possiblity to remove correlation maxima
; for a second linear regression round

    XMessage, ['Use mouse to remove maxima:', $
               '  - Left button to select new area', $
               '  - Right button to terminate selection', $
               '', $
               '- Middle button to change spline parameter', $
               '  (number of points to add)'], $
      TITLE='Mouse intput', XOFF=10, WBASE=wbase
    
    removeMaxima = -1
    REPEAT BEGIN
        Cursor, x0, y0, /DATA, /DOWN
        Cursor, x1, y1, /DATA, /UP
        IF !MOUSE.BUTTON EQ 1 THEN BEGIN ; left button was pressed
            IF x0 GT x1 THEN SwapVars, x0, x1
            IF y0 GT y1 THEN SwapVars, y0, y1
            OPLOT, [x0,x1,x1,x0,x0], [y0,y0,y1,y1,y0]
            ; convert clicked coords to array indices
            IF KEYWORD_SET(LOGFREQ) THEN BEGIN
                highInd = WHERE(newYAxis[yMin:yMax] GE y0)
                highInd = highInd[N_ELEMENTS(highInd) - 1]
                lowInd = (WHERE(newYAxis[yMin:yMax] LE y1))[0]
                IF lowInd EQ -1 THEN highInd = -1
            ENDIF ELSE BEGIN
                highInd = ROUND((y0-newYAxis[yMin])/yScale)
                lowInd = ROUND((y1-newYAxis[yMin])/yScale)
            ENDELSE
            IF (lowInd GT (yMax-yMin)) OR (highInd LT 0) OR $
              (lowInd GT highInd) THEN doEliminate = 0 $ 
            ELSE BEGIN
                doEliminate = 1
                lowInd = lowInd>0
                highInd = highInd<(yMax-yMin)
                indInterval = LINDGEN(highInd-lowInd+1)+lowInd
            ENDELSE
           
            IF doEliminate EQ 1 THEN BEGIN
            ; search for times that lie in the specified interval 
                maximaInd = WHERE((newCcXAxis[maxima[lowInd:highInd]] LE x1) $
                                  AND (newCcXAxis[maxima[lowInd:highInd]] GE x0), $
                                  numRemv)
                IF numRemv GT 0 THEN BEGIN
                    IF KEYWORD_SET(LOGFREQ) THEN BEGIN
                        OPLOT, newCcXAxis[(maxima[lowInd:highInd])[maximaInd]], $
                          ((newYAxis[ymin:ymax])[lowInd:highInd])[maximaInd], $
                          COLOR=0, PSYM=6
                    ENDIF ELSE BEGIN
                        OPLOT, newCcXAxis[(maxima[lowInd:highInd])[maximaInd]], $
                          newYAxis[yMin]+yScale*indInterval[maximaInd], $
                          COLOR=0, PSYM=6
                    ENDELSE
                    IF (removeMaxima[0] GE 0) THEN $
                      removeMaxima = [removeMaxima, indInterval[maximaInd]] $
                    ELSE removeMaxima = indInterval[maximaInd]
                ENDIF
            ENDIF
        ENDIF ELSE IF !MOUSE.BUTTON EQ 2 THEN BEGIN ; middle button was pressed
            PRINT
            PRINT, 'Change spline interpolation parameter.'
            Read_Test, 'Enter number of points to add between original data points', $
              addPts, newVal, MINIMUM = 0.0
            IF ROUND(newVal) EQ addPts THEN PRINT,  '  -> Not changed' $
            ELSE BEGIN
                PRINT
                XKill, wbase
                addPts = FLOOR(newVal)
                isSecondRound = 1
                GOTO,DO_SPLINE
            ENDELSE
        ENDIF
    ENDREP UNTIL !MOUSE.BUTTON EQ 4 ; until right button is pressed
    XKill, wbase

    IF KEYWORD_SET(PSOUT) THEN BEGIN
        entry_device = !D.NAME
        SET_PLOT, 'PS'
        DEVICE, FILENAME='drift_rate_1.ps', /COLOR, BITS_PER_PIXEL=8, $
          /PALATINO, /BOLD, /LANDSCAPE
        Drift_Plot, ccImage, ccXAxis, newCcXAxis, newYAxis[yMin:yMax], yScale, $
          maxima, INDGEN(N_ELEMENTS(maxima)), const, slope[0], LOGFREQ=LOGFREQ
        DEVICE, /CLOSE
        SET_PLOT, entry_device
        PRINT
        PRINT, "Saved first plot to postscript file 'drift_plot_1.ps'."
        PRINT
    END
    
    IF (removeMaxima[0] EQ -1) OR (N_ELEMENTS(removeMaxima) GE N_ELEMENTS(maxima)) THEN BEGIN
        PRINT, '  No correlation maxima selected for removal or all maxima removed. '
        PRINT, '  Regression is not recalculated.'
        PRINT, '--------------------------------------------------------------------'
        PRINT
        !P.TITLE = oldPTitle
        !X.TITLE = oldXTitle
        !Y.TITLE = oldYTitle
        WSET, oldWinNr
        RETURN
    ENDIF
    
    keepMaxima = maxima
    keepMaxima[removeMaxima] = -1
    keepMaxima = WHERE(keepMaxima GE 0)

; calculate new regression, plot results in new window
    
    slope = REGRESS(TRANSPOSE((newYAxis[yMin:yMax])[keepMaxima]), $
                    newCcXAxis[maxima[keepMaxima]], $
                    REPLICATE(1.0, N_ELEMENTS(keepMaxima)), $
                    xfit, const, sigma, ftest, corr, rmul, chisq, /RELATIVE_WEIGHT)
    
    winNr = !D.WINDOW + 1
    WINDOW, winNr, TITLE = 'Linear regression (2nd iteration)'

    Drift_Plot, ccImage, ccXAxis, newCcXAxis, newYAxis[yMin:yMax], yScale, $
      maxima, keepMaxima, const, slope[0], LOGFREQ=LOGFREQ

    PRINT
    PRINT, '  ', STRN(N_ELEMENTS(removeMaxima)), $
      ' correlation maxima selected for removal. '
    PRINT, '  Recalculating regression. '
    
    PRINT
    PRINT, '  Linear regression - second iteration:'
    PRINT, '	Sigma: ', STRN(sigma)
    PRINT, '	Corr: ', STRN(corr)
    PRINT, '	Chisq: ', STRN(chisq)
    PRINT
    IF KEYWORD_SET(LOGFREQ) THEN $
      PRINT, '  Slope: ',STRN(1./slope[0]),' +/- ',STRN(sigma/(slope[0]^2)), $
        ' [1/s]' $
    ELSE $
      PRINT, '  Drift rate: ',STRN(1./slope[0]),' +/- ',STRN(sigma/(slope[0]^2)), $
      ' [MHz/s] at ',centerFreq,' MHz'

    IF KEYWORD_SET(PSOUT) THEN BEGIN
        entry_device = !D.NAME
        SET_PLOT, 'PS'
        DEVICE, FILENAME='drift_rate_2.ps', /COLOR, BITS_PER_PIXEL=8, $
          /PALATINO, /BOLD, /LANDSCAPE
        Drift_Plot, ccImage, ccXAxis, newCcXAxis, newYAxis[yMin:yMax], yScale, $
          maxima, keepMaxima, const, slope[0], LOGFREQ=LOGFREQ
        DEVICE, /CLOSE
        SET_PLOT, entry_device
        PRINT
        PRINT
        PRINT, "Saved second plot to postscript file 'drift_plot_2.ps'."
    END
    
    PRINT
    PRINT, '--------------------------------------------------------------------'
    PRINT

    WSET, oldWinNr    
    !P.TITLE = oldPTitle
    !X.TITLE = oldXTitle
    !Y.TITLE = oldYTitle
END