;+
; NAME:
;	PH_DITHER
;
; PURPOSE:
;	Reduction of dithered/mosaic PHARO images.
;
; EXPLANATION:
;	Determines telescope pointing from image FITS headers, then ajusts
;	relative positions of images by cross-correlation of overlap regions.
;	Saves both a "mosaic" average of all frames including non-overlap
;	areas, and a "mean image" of only the central overlap area.
;
; CALLING SEQUENCE:
;	PH_DITHER, mosaic, KWD=KWD, ...
;
; INPUTS:
;       objfr=  Vector of frame numbers for object.  If not given, will
;               prompt for first and last frame numbers.
;
;       skyfr=  Vector of frame numbers for sky frames.  If not given, will
;               use "off" pointings in object frames.
;
;	flatname= Filename of flat to use.  If not specified, uses
;		'flat_X_YYmas.fits' where X is filter and YY is platescale.
;
;	bpixname= Filename of bad pixel map to use.  If not specified, uses
;		'bpix_X_YYmas.fits' where X is filter and YY is platescale.
;
;       datadir= Full path to data directory.  If not specified, determined
;               from date.
;
;       procdir= Full path to processing directory.  If not specified,
;               determined from date.
;
;       caldir= Full path to calibration file directory.  If not specified,
;               assumed to be same as processing directory
;
;       date=   UT date in 'YYMMDD' format.  Used only if datadir and
;               procdir are not defined.
;
;	mask=	1024x1024 pixel mask image (or weight) to give image image
;		in the mosaic.  Use this to mask out bad regions of the
;		detector.
;
;	/autocray Auto-filter cosmic rays (watch out - can chop top off PSF!)
;
;       /noshift Do no cross-correlate and shift images in final stack.
;
;	skyrej=	Fraction of highest pixel values to reject from median stack.
;		Recommend using 1/(N-1) where N is number of dither positions.
;		Default = 0.0
;
;	/fixbanding Attempt to remove banding in images due to cross-talk
;		on bright objects.
;
;       /silent Supress informational messages
;
;	mag=	1/pixel step size for Fourier shift during coalignment of
;		frames.  Significantly increases processing time.  Default=1
;
; OUTPUTS:
;	mosaic = mosaic image (NxNx2) where N is larger than the size
;		of input frames, and the layers are...
;		0	Average of sky-subtracted frames.
;		1	Number of frames which contributed to each pixel
;
;	mean = mean image (NxNx2) where N is the size of input frames.
;
;	Writes to directory 'procdir' the following files:
;		ph????r1.fits	Reduced individual frames.
;		ph????r2m.fits	Final mosaic image.
;		ph????r2d.fits	Mean image, trimmed to original pixel size.
;
; MODIFICATION HISTORY:
;       Original, based on NIRC2DITH, written 05/07
;					A. Bouchez, Caltech Optical Obs.
;	Replaced UNSTACK call with RFITS, 5/22/07, A. Bouchez, COO.
;	Dithered output (_r2d) now same size as input frames; added
;	  /fixbanding keyword, reduced output array size. 4/9/09, AHB.
;-
PRO     PH_DITHER, im4, im5, im6, im7, objfr=objfr, skyfr=skyfr, $
          flatname=flatname, bpixname=bpixname, datadir=datadir, $
          procdir=procdir, caldir=caldir, date=date, mask=mask, $
          autocray=autocray, noshift=noshift, skyrej=skyrej, $
          fixbanding=fixbanding, silent=silent, mag=mag

;;; Define constants

  fmt = '("ph",I4.4,".fits.gz")'        ; PHARO filename format
  fmt_r1 = '("ph",I4.4,"r1.fits")'	; Filename format for saved frames
  fmt_r2m = '("ph",I4.4,"r2m.fits")'	; Filename format for saved mosaic
  fmt_r2d = '("ph",I4.4,"r2d.fits")'	; Filename format for saved coadd image
  cr_zpt = 335.8                        ; Cass ring zeropoint (N up)
  drd = 2.0				; Dither position radius (")
  csig = 10.0				; Cosmic ray filter threshold
  xcor_min = 5.0			; Min x-correl. to trust
  xshi_max = 100.0			; Max offset from pred. dither to allow
  im_pad = 200				; Pixel padding of mosaic image

;;; Parse keywords

  if not KEYWORD_SET(noshift) then noshift = 0
  if not KEYWORD_SET(silent) then silent = 0
  if not KEYWORD_SET(autocray) then autocray = 0
  if not KEYWORD_SET(skyrej) then skyrej = 0
  if not KEYWORD_SET(mag) then mag = 1.0	; x-correlation magnification

  if not (KEYWORD_SET(datadir) and KEYWORD_SET(procdir)) then begin
    if not KEYWORD_SET(date) then begin
      date = ''
      READ, 'Please enter date ("YYMMDD"): ', date
    endif
    date = STRING(date,f='(I6.6)')
    datadir = '/data/pharo_' + date + '/'
    procdir = '/Users/abouchez/pharo/' + date + '/'
  endif

  if not KEYWORD_SET(caldir) then caldir = procdir

  if not KEYWORD_SET(objfr) then begin
    fstr= ''
    READ, 'Please enter OBJECT first frame number: ', fstr
    lstr = ''
    READ, 'Please enter OBJECT last frame number: ', lstr
    flno = FIX([fstr, lstr])
    objfr = flno[0] + LINDGEN(flno[1]-flno[0]+1)
  endif

  if not KEYWORD_SET(skyfr) then begin
    if not silent then $
      MESSAGE, /info, 'Using dithered object frames for sky.'
    skyfr=objfr
    dsky = 1B
  endif else dsky = 0B

;;; Determine image size, filter, platscale, and PA from first image header

  hd0 = HEADFITS(datadir + STRING(objfr[0],f=fmt), /silent)
  sz0 = [2*SXPAR(hd0,'NAXIS1'), 2*SXPAR(hd0,'NAXIS2'), N_ELEMENTS(objfr)]

  filt = STRTRIM(SXPAR(hd0,'FILTER'),2)		; USED ONLY TO CHOOSE FLAT
  case filt of
    'K_short': filt='Ks'
    'Br-gamma': filt='Ks'
    else: foo=0					; Add more lines as necessary!
  endcase

  camname = STRTRIM(SXPAR(hd0,'CAROUSEL'),2)
  case camname of
    '25 mas': pscl = 0.025
    '40 mas': pscl = 0.040
  endcase

  pa = (cr_zpt - SXPAR(hd0,'CR_ANGLE') + 360) mod 360		; PA of N

;;; Read flat and bad pixel maps (does not yet deal correctly with subarrays!)

  if not KEYWORD_SET(flatname) then $
    flatname = 'flat_' + filt + '_' + $
      STRING(ROUND(pscl*1e3),f='(I0.0)') + 'mas.fits'
  tmp = FILE_SEARCH(caldir + flatname)
  while (STRLEN(tmp[0]) eq 0) do begin
    MESSAGE, /info, 'Flat file ' + flatname + ' not found.'
    READ, 'Please enter a new file name: ', flatname
    tmp = FILE_SEARCH(caldir + flatname)
  endwhile
  if not silent then $
    MESSAGE, /info, 'Using flat: ' + flatname
  flat = READFITS(caldir + flatname, /silent)

  if not KEYWORD_SET(bpixname) then $
    bpixname = 'bpix_' + filt + '_' + $
      STRING(ROUND(pscl*1e3),f='(I0.0)') + 'mas.fits'
  tmp = FILE_SEARCH(caldir + bpixname)
  while (STRLEN(tmp[0]) eq 0) do begin
    MESSAGE, /info, 'Bad pixel file ' + bpixname + ' not found.'
    READ, 'Please enter a new file name: ', bpixname
    tmp = FILE_SEARCH(caldir + bpixname)
  endwhile
  if not silent then $
    MESSAGE, /info, 'Using bpix: ' + bpixname
  bpix = READFITS(caldir + bpixname, /silent)

;;; Sort all object frames into distinct pointings based on RA/Dec

  nfr = N_ELEMENTS(objfr)
  hd = STRARR(N_ELEMENTS(hd0), nfr)
  rade = DBLARR(2, nfr)
  for n=0,nfr-1 do begin
    hd[*,n] = HEADFITS(datadir + STRING(objfr[n],f=fmt), /silent)
    rade[*,n] = [SXPAR(hd[*,n],'CRVAL1'), SXPAR(hd[*,n],'CRVAL2')]	; deg
  endfor

  mrade = [TOTAL(MINMAX(rade[0,*])), $			; Midpoint RA/Dec (deg)
           TOTAL(MINMAX(rade[1,*]))] / 2.
  rdoff = (rade - REBIN(mrade,2,nfr)) * $		; Offset RA/Dec (")
    [COS(rade[1,*]*!dtor), REPLICATE(1.,1,nfr)] * 3600
  dgroup = ROUND(rdoff[0,*]/drd) + 100L * ROUND(rdoff[1,*]/drd)
  duniq = dgroup(UNIQ(dgroup))
  dpos = N_ELEMENTS(duniq)
  if not silent then $
    MESSAGE, /info, 'Found ' + STRING(dpos,f='(I0.0)') + $
    ' distinct object pointings'
  didx = [REFORM(objfr,1,nfr), LONARR(1,nfr)]		; [frno, pointing]
  for i=0,dpos-1 do $
    didx[1, WHERE(dgroup eq duniq[i])] = i

;;; Compute median sky frame for each pointing

  if dsky then begin				; if using dithered sky frames
    sky = LONARR(sz0[0], sz0[1], dpos)
    skydat = LONARR(3, dpos)
    for n=0,dpos-1 do begin
      offfr = objfr[WHERE(didx[1,*] ne n)]
      skydat[*,n] = [N_ELEMENTS(offfr), MINMAX(offfr)]
      if not silent then $
        MESSAGE, /info, 'Computing sky ' + $
        STRING([n+1,dpos],f='(I0.0," of ",I0.0)') + ' using ' + $
        STRING(skydat[0,n],f='(I0.0)') + ' frames.'
      sky[*,*,n] = MEDIM(datadir + STRING(offfr,f=fmt), /pharo, reject=skyrej)
    endfor

  endif else begin				; if using indep. sky frames
    nskyfr = N_ELEMENTS(skyfr)
    if not silent then $
      MESSAGE, /info, 'Computing median of ' + $
      STRING(nskyfr,f='(I0.0)') + ' independent sky frames.'
    onesky = MEDIM(datadir + STRING(skyfr,f=fmt), /pharo, reject=skyrej)
    sky = REBIN(onesky, sz0[0], sz0[1], dpos)
    skydat = REBIN([nskyfr, MINMAX(skyfr)], 3, dpos)
  endelse

;;; Determine pixel offsets of obj frames, set up image arrays, mosaic weights

  xyoff = FLTARR(2,sz0[2])
  rot = [[   COS(pa*!dtor), SIN(pa*!dtor)], $
         [-1*SIN(pa*!dtor), COS(pa*!dtor)]]
  for n=0,sz0[2]-1 do $
    xyoff[*,n] = (([-1,1]*rdoff[*,WHERE(didx[0,*] eq objfr[n])]) # rot)/pscl
  sz1 = sz0[0:1] + ROUND([MAX(xyoff[0,*])-MIN(xyoff[0,*]), $	;Pad...
    MAX(xyoff[1,*])-MIN(xyoff[1,*])]) + im_pad
  for i=0,1 do $                                                ;Make axes even
    if sz1[i]-2*(sz1[i]/2) eq 1 then sz1[i]=sz1[i]+1

  if not silent then $
    MESSAGE, /info, 'Creating image arrays (up to ' + $
    STRING([sz1,sz0[2]],f='(I0.0,"x",I0.0,"x2x",I0.0)') + ')'
  im4 = FLTARR(sz0[0],sz0[1],sz0[2])
  im5 = FLTARR(sz1[0],sz1[1],2,sz0[2])		; 0=img, 2=pixel weight
  im6 = FLTARR(sz1[0],sz1[1],2)			; same
  xx = REBIN(INDGEN(sz1[0],1),sz1[0],sz1[1])
  yy = REBIN(INDGEN(1,sz1[1]),sz1[0],sz1[1])

  if not KEYWORD_SET(mask) then begin
    wt = 1						; rel. weight per pixel
    ;DIST_CIRCLE,r,sz0[0],sz0[0]/2-0.5,sz0[1]/2-0.5	; dist from img center
    ;wt = 1-(SQRT(2)*r/sz0[0])^2.0
  endif else wt = mask

  c0 = [sz1, sz1]*0.5 + [-0.5*sz0[0:1], 0.5*sz0[0:1]-1]

;;; Reduce individual object frames.

  for n=0,sz0[2]-1 do begin
    im0 = RFITS(datadir + STRING(objfr[n],f=fmt), hd, /pharo, /silent)
    idx = didx[1,(WHERE(didx[0,*] eq objfr[n]))[0]]
    im1 = (im0 - sky[*,*,idx]) / flat
    im2 = BPIXFIX(im1, bpix, neighbors=3)
    if autocray then $
      im3 = CRAYFIX(im2, t=csig) else $
      im3 = im2

    SXADDPAR, hd, 'BITPIX', -32
    SXADDPAR, hd, 'NAXIS', 2
    SXADDPAR, hd, 'NAXIS1', sz0[0]
    SXADDPAR, hd, 'NAXIS2', sz0[1]
    SXDELPAR, hd, 'NAXIS3'
    SXADDPAR, hd, 'IMGLEVEL', 1, ' Image processing level (pre-reduced)'
    SXADDPAR, hd, 'IMGTYPE', 'frame', ' Image type (frame/mosaic/mean)'
    SXADDPAR, hd, 'DITHPOS', idx, ' Dither position'
    SXADDPAR, hd, 'SKYNFR', skydat[0,idx], ' Number of sky frames'
    SXADDPAR, hd, 'SKYFRNO0', skydat[1,idx], ' First sky frame'
    SXADDPAR, hd, 'SKYFRNO1', skydat[2,idx], ' Last sky frame'
    SXADDPAR, hd, 'FLATNAME', flatname
    SXADDPAR, hd, 'BPIXNAME', bpixname
    if autocray then $
      SXADDPAR, hd, 'CRTHRESH', csig, ' Cosmic ray threshold (sigma)' else $
      SXADDPAR, hd, 'CRTHRESH', 0, ' Cosmic ray filtering disabled'
    sname = STRING(objfr[n],f=fmt_r1)
    if not silent then MESSAGE, /info, 'Saving ' + sname
    WRITEFITS, procdir + sname, im3, hd
    if n eq 0 then hd0 = hd

;;; Insert at estimated location in mosaic, subtract residual sky or banding

    im4[*,*,n] = im3			; NO DISTORTION CORRECTION YET!

    if KEYWORD_SET(fixbanding) then begin	; To subtract cross-talk,
      im4sky = FLTARR(sz0[0], sz0[1])		;  subtract median of each row
      for j=0,sz0[1]-1 do begin
        rpix = im4[*,j,n]
        im4sky[*,j] = MEDIAN( (rpix[SORT(rpix)])[0:sz0[0]*0.50-1] )
      endfor
    endif else $
      SKY, im4[*,*,n], im4sky, /silent		; Otherwise, use median sky val

    c1 = c0 + [xyoff[*,n], xyoff[*,n]]
    if n eq 0 then begin
      im5[c1[0]:c1[2],c1[1]:c1[3],0,0] = im4[*,*,0] - im4sky
      im5[c1[0]:c1[2],c1[1]:c1[3],1,0] = wt
    endif else begin
      tmp = FLTARR(sz1[0],sz1[1],2)
      tmp[c1[0]:c1[2],c1[1]:c1[3],0] = im4[*,*,n] - im4sky
      tmp[c1[0]:c1[2],c1[1]:c1[3],1] = wt
      olap = WHERE(im6[*,*,1] and tmp[*,*,1])

;;; 7. If overlap, cross-correlate over overlap region, correct, and add.

      if (olap[0] ne -1) and (not KEYWORD_SET(noshift)) then begin
        c2 = [MIN(xx[olap]), MIN(yy[olap]), MAX(xx[olap]), MAX(yy[olap])]
        if (c2[2]-c2[0]+1)/2*2 ne c2[2]-c2[0]+1 then c2[2]=c2[2]-1	;make
        if (c2[3]-c2[1]+1)/2*2 ne c2[3]-c2[1]+1 then c2[3]=c2[3]-1	; even
        im6a = im6[c2[0]:c2[2],c2[1]:c2[3],0]
        im6b = tmp[c2[0]:c2[2],c2[1]:c2[3],0]
        sz02 = [c2[2]-c2[0], c2[3]-c2[1]] + 1

        u = [FINDGEN(sz02[0]*mag/2),FINDGEN(sz02[0]*mag/2)-(sz02[0]*mag/2)]
        v = [[FINDGEN(1,sz02[1]*mag/2)], $
             [FINDGEN(1,sz02[1]*mag/2)-(sz02[1]*mag/2)]]
        uu = REBIN(u,sz02[0]*mag,sz02[1]*mag)/mag
        vv = REBIN(v,sz02[0]*mag,sz02[1]*mag)/mag
        cc = FFT( PFFT(im6a,mag) * CONJ(PFFT(im6b,mag)), 1)
        maxcc = MAX(cc, mp)

        w = WHERE(FLOAT(cc)-MEDIAN(FLOAT(cc)) lt 3*STDDEV(FLOAT(cc)))
        if (w[0] ne -1) then $
          mccsig = (FLOAT(maxcc) - MEDIAN(float(cc))) / STDDEV(FLOAT(cc[w])) $
          else mccsig = -1
        mshift = MAX(ABS([uu[mp],vv[mp]]))

        if (mccsig lt xcor_min) or (mshift gt xshi_max) then $
          c2 = c1 else $
          c2 = c1 + [uu[mp], vv[mp], uu[mp], vv[mp]]
      endif else begin
        c2 = c1
        if not silent then $
          MESSAGE, /info, 'No overlap. No shift.'
        uu = 0.
        vv = 0.
        mp = 0
      endelse
        
      im5[c2[0]:c2[2],c2[1]:c2[3],0,n] = $
        FSHIFT(im4[*,*,n]-im4sky,uu[mp]-FIX(uu[mp]),vv[mp]-FIX(vv[mp]))
      im5[c2[0]:c2[2],c2[1]:c2[3],1,n] = wt
    endelse

;;; Computed and save weighted average over all images in stack

    im6[*,*,1] = REBIN(im5[*,*,1,0:n], sz1[0], sz1[1])	; normalized weight
    wt6 = im6[*,*,1]
    wt6(WHERE(wt6 eq 0)) = 1.0
    im6[*,*,0] = REBIN(im5[*,*,0,0:n]*im5[*,*,1,0:n], sz1[0], sz1[1]) / wt6
  endfor

  SXADDPAR, hd0, 'NAXIS', 3
  SXADDPAR, hd0, 'NAXIS1', sz1[0]
  SXADDPAR, hd0, 'NAXIS2', sz1[1]
  SXADDPAR, hd0, 'NAXIS3', 2, after='NAXIS2'
  SXADDPAR, hd0, 'IMGLEVEL', 2, ' Image processing level (mean image)'
  SXADDPAR, hd0, 'IMGTYPE', 'mosaic', ' Image type (frame/mosaic/mean)'
  SXADDPAR, hd0, 'OBSNFR', sz0[2], ' Frames in dither pattern'
  SXADDPAR, hd0, 'FRNO0', MIN(objfr), ' First object frame'
  SXADDPAR, hd0, 'FRNO1', MAX(objfr), ' Last object frame'
  sname = STRING(objfr[0], f=fmt_r2m)
  if not silent then $
    MESSAGE, /info, 'Saving ' + STRING(sz1,f='(I0.0,"x",I0.0)') + $
      ' mosaic as ' + sname
  WRITEFITS, procdir + sname, im6, hd0

;;; Save weighted average of central region (same size as input frames)

  c3 = [sz1[0]/2-sz0[0]/2, sz1[1]/2-sz0[1]/2, $
        sz1[0]/2+sz0[0]/2-1, sz1[1]/2+sz0[1]/2-1]
  im7 = im6[c3[0]:c3[2],c3[1]:c3[3],*]

  SXADDPAR, hd0, 'NAXIS', 2
  SXADDPAR, hd0, 'NAXIS1', sz0[0]
  SXADDPAR, hd0, 'NAXIS2', sz0[1]
  SXADDPAR, hd0, 'IMGTYPE', 'mean', ' Image type (frame/mosaic/mean)'
  sname = STRING(objfr[0], f=fmt_r2d)
  if not silent then $
    MESSAGE, /info, 'Saving ' + STRING(sz0[0:1],f='(I0.0,"x",I0.0)') + $
      ' mean image as ' + sname
  WRITEFITS, procdir + sname, im7, hd0

END