;+
; NAME: shift_data.pro
;
; PURPOSE: Coaligns image sequences with frame-to-frame shifts to
;  sub-pixel accuracy.  
;
;  By default, the shifts are assumed unknown, and are calculated by
;  Fourier cross correlation.  If shifts have been previously
;  determined, set /RERUN and input the (2, nt-1) array of shifts.
;
;  Data are then shifted (in the *OPPOSITE* direction of the
;  frame-to-frame displacement!) to sub-pixel accuracy using Fourier
;  interpolation. Finally, "wrapped" data are zeroed.
;
; USAGE: shifted = shift_data(unshifted, shifts=shifts,
;                             init=init, rerun=rerun)
;
; INPUT: Data cube of images, fltarr(nx,ny,nt).
;     
; OPTIONAL INPUT: SHIFTS = In case shifts have been previously
;  determined, set /RERUN and input the (2, nt-1) array of shifts.
;
; OUTPUT: Data shifted to coalign to sub-pixel accuracy; optionally
;    returns shifts.
;
; KEYWORDS:
;   INIT = by default, shifts n-th frame to match init. frame;
;          set to 0 use frame-to-frame matching
;   RERUN = set to 1 if pre-determined shifts are to be applied
;           (e.g., use B_z to find shifts, then coalign B_x, B_y
;            without recalculating shifts)
;   QUIET = set to suppress printing of progress.
;
; HISTORY: Cleaned up an earlier version, now uses IDL routines
;   CROSS_COR_TAYLOR.PRO and SHIFT_FRAC2D. - 2010/12/15, BTW
;   
;-
function shift_data, unshifted, shifts=shifts, init=init, $
                     rerun=rerun, quiet=quiet 

nx = n_elements(unshifted(*,0,0))
ny = n_elements(unshifted(0,*,0))
nt = n_elements(unshifted(0,0,*))


; The default is to shift position to match the initial image.
if (n_elements(init) eq 0) then init = 1.

; If RERUN=0, then calculate shifts; otherwise, re-use old shifts.
if not(keyword_set(rerun)) then shifts = fltarr(2,nt-1)
if (n_elements(shifts(*,0)) ne 2 or $
    n_elements(shifts(0,*)) ne (nt-1) or $
    n_elements(shifts) ne 2*(nt-1)) then begin
   message,'Incorrect number of shifts for rerunning.',/info
   stop
endif

; Output (shifted) array will be fltarr(nx,ny,nt) 
shifted = fltarr(nx,ny,nt)
shifted(*,*,0) = unshifted(*,*,0) ; first slice of the output

; Avoid edge effects by co-registering central 80%
xedge_pad = round(0.1*nx)
xi = xedge_pad
xf = nx - xedge_pad
yedge_pad = round(0.1*ny)
yi = yedge_pad
yf = ny - yedge_pad

template = shifted(xi:xf,yi:yf) 
template = template - mean(template) ; Make sure to demean the data!

; Print progress
if not(keyword_set(quiet)) then print, $
   'Step Number,     X-shift,     Y-shift:'

for i=1.,nt-1. do begin

   if not(keyword_set(rerun)) then begin ; don't recalc shifts if rerun
      iarr = unshifted(xi:xf,yi:yf,i)
      iarr = iarr - mean(iarr)  ; always demeaning!
   
      cross_cor_taylor, template, iarr, ccor, $
                        shiftx, shifty, /quiet

   endif else begin ; if rerun, re-use old shifts
      shiftx = shifts(0, i-1)
      shifty = shifts(1, i-1)
   endelse

   ; To counteract inferred shift, apply *negative* shift
   shifted_i = shift_frac2d(unshifted(*,*,i), -shiftx, -shifty)

   ; Now zero out "wrapped" data
   if (shiftx gt 0) then begin
      dx_crop = ceil(shiftx)
      xcrop_i = nx-dx_crop
      xcrop_f = nx-1
   endif else begin
      dx_crop = abs(floor(shiftx))
      xcrop_i = 0
      xcrop_f = dx_crop
   endelse
   shifted_i(xcrop_i:xcrop_f,*)=0

   if (shifty gt 0) then begin
      dy_crop = ceil(shifty)
      ycrop_i = ny-dy_crop
      ycrop_f = ny-1
   endif else begin
      dy_crop = abs(floor(shifty))
      ycrop_i = 0
      ycrop_f = dy_crop
   endelse
   shifted_i(*, ycrop_i:ycrop_f)=0

   ; Add shifted data from i-th step to array of all shifted data
   shifted(*,*,i) = shifted_i

   ; If calculating shifts, store them, and make new template
   if not(keyword_set(rerun)) then begin

      shifts(*,i-1) = [temporary(shiftx), temporary(shifty)]

      ; Match current image with initial image, or the previous image?
      if (keyword_set(init)) then begin 
         template = shifted_i(xi:xf,yi:yf) ; match to *shifted* previous
         template = template - mean(template) ; still demeaning!
      endif else begin
         template = shifted_i(xi:xf,yi:yf) ; match to *unshifted* previous
         template = template - mean(template) ; still demeaning!
      endelse
   endif 

   ; Print progress
   if (not(keyword_set(quiet)) and ((i-1)mod(20) eq 0)) then $
      print,i-1,shifts(*,i-1)
      
endfor

return, shifted

end