PRO xrt_spotcor,index,data,index_out,data_out,numspots,spotmap, $
        thresh=thresh,all=all,abort=abort,quiet=quiet
; =========================================================================
;+   
;
; PROJECT:
;       Solar-B / XRT
;
; NAME:
;
;       XRT_SPOTCOR
;
; CATEGORY:
;
;       Hinode XRT data analysis.
;
; PURPOSE:
;
;       Cosmetic correction for CCD spots and dust on XRT data.
;       Uses thin-plate splines for an improved result.
;
; CALLING SEQUENCE:
;
;      XRT_SPOTCOR,index,data,index_out,data_out,[numspots],[spotmap], $
;         [thresh=thresh], [all=all], [abort=abort], [quiet=quiet]
;
; INPUTS:
;
;       INDEX        - [Mandatory] (a standard XRT index format) 
;                      Index structure for the data.
;                      Note: this routine is intended for Level 1 data (i.e.
;                      data that has been prepared by xrt_prep - this routine
;                      is called within xrt_prep if despotting is requested).
;
;       DATA         - [Mandatory] (Nx x Ny x N_images) 
;                      Data cube of consisting of images of the same
;                      size and binning and ccd location.
;
; KEYWORDS:
;
;       TRESH        - [Optional] 
;                       Threshold of fraction contamination by spots above a
;                       pixel is corrected [default=0.5]. It may be unnecessary
;                       to correct an 4x4 binned pixel for your purposes, if, 
;                       e.g., only one of the 1x1 pixels which are binned to 
;                       form it is a spot.  In this case the user might want 
;                       to set thresh=1/16.
;        ALL         - [Optional] Try to fix *all* spots (without testing 
;                        for their depth)
;        ABORT       - [Optional] Boolean indicating if routine
;                        completed successfully (abort=0).   
;        QUIET       - [Optional] Boolean indicating if routine
;                        runs without printing (quiet=1).   
;
; OUTPUTS:
;
;       DATA_OUT       - [Mandatory] (image array Nx x Ny x N_images)
;                      Spot and dust corrected XRT images.
;       INDEX_OUT      - [Mandatory] (standard XRT index format)
;                      Modified index structure for the data.  
;       NUMSPOTS       - [Optional] (8 x N_images floating array)
;                      Diagnostics of spots in each filter k, where
;                      for NUMSPOTS(j,k):
;                  j=0: number of original spot pixels (n_spotpix) in full 1x1
;                    1: n_spotpix pixels in 1x1 after selection by filter
;                         ratio (currently same as j=0)
;                    2: n_spotpix in 1x1 after expansion
;                    3: n_spotpix in Nx x Ny FOV after expansion & thresholding
;                    4: number of distinct spots in Nx x Ny FOV
;                    5: n_spotpix in Nx x Ny FOV without expansion
;                         (not calculated in all cases)
;                    6:  final number of distinct spots corrected in filter
;                    7:  final n_spotpix corrected in filter
;        SPOT_MAP     - [Optional]  Nx x Ny x N_filter array
;                     Binary map of corrected spots + "dust growth" for 
;                     each filter
;                      
; EXAMPLES:
;
;       Basic call to get a corrected array of image
;        IDL> xrt_spotcor, index, data, index_out, data_out
;
;       Call to get a corrected array of image, spot arrays, statistics, 
;         and change spot threshold from default (=0.5):
;      IDL> xrt_spotcor, index, data, index_out, data_out, numspots, spotmap, $
;                thresh=thresh
;
;
; COMMON BLOCKS:
;
;       None
;
; NOTES:
;
;       Make sure your arrays are composed of images all the same size,
;       binning, CCD location, and spot epoch. If they are not an error 
;       message will appear and you will be ejected from the program.
;
;       The routine does not work on Gband data. 
;
;       It does not work well for 8x8 binned data (too heavily binned, spots
;       too blurred together to work well). For 4x4 data, not all spots
;       are corrected equally well.
;    
;       The program assumes the timespan of the data in the cube is short 
;       enough that evolution of the underlying structures is not 
;       important.  It defines the spots which require correction based
;       on the first image in each filter. Since the effect of the spots
;       depends on the temperature of the underlying feature, their
;       evolution over the timespan of the cube should be minimal for
;       the spot correction to be optimal.
; 
;       The correction is not photometric, though should be reasonably
;       good in most cases, provided the above guidelines are followed.
;
;       The latest versions (after the August 2013 update) include a cosmetic
;       correction for dust on the CCD and the growth seen in the dust after
;       bakeouts post light leak (May 2012).  The latter "dust growth" pixels
;       are added to the output spot array.  The model for the "dust growth"
;       will likely be improved over time. 
;       
;
; MODIFICATION HISTORY:
;
;      17-Apr-2009 - S.Saar and A. Savcheva  (written)
;       5-May-2009 - S. Saar (some final additions & adjustments)
;      27-May-2009 - S. Saar (changed location of spot data files)
;       5-Jun-2009 - S. Saar (corrected/generalized path to spot data files)
;       1-Feb-2013 - S. Saar (modified correction for 4x4 binned data, added
;                             note above, standardized retrieval of spot data
;                             with <xrt_tup_contam>, changed gotos to returns,
;                             added abort, quiet, history update to index)
;      18-Aug-2013 - S. Saar (set new default threshold to 0.5, added thresh
;                             to history update, now also makes a cosmetic dust
;                             correction and an approximate correction for the 
;                             increased size of dust between bakeouts, post 
;                             light leak. Updated history to include dust info.
;                             Much helpful testing by P. McCauley anf J. Slavin)
;      23-Sep-2013 - S. Saar (can now do cal frame images, rejects G bands and
;                             darks. added ALL keyword. Changed how dust growth
;                             is treated (now using xrt_dustgrow).)
;      21-Feb-2014 - S. Saar (fixes crash for odd images with fw1=fw2=0)
;      14-Apr-2014 - S. Saar (now IDs spots on the edge of the FOV) 
;      01-Mar-2017 - J. Slavin - Added fix for case when missing frame pixels
;                             (value = -999.) are in spot region.  Also, made
;                             sure that those pixels are not filled in as a
;                             result of spot correction (left as -999.)
;
;- 
; =========================================================================
progver = 'v2009-Jun-05' ;--- (S. Saar) 
progver = 'v2013-Feb-01' ;--- (S. Saar) 
progver = 'v2013-Aug-18' ;--- (S. Saar) 
progver = 'v2013-Sep-23' ;--- (S. Saar) 
progver = 'v2014-Feb-21' ;--- (S. Saar) 
progver = 'v2014-Apr-14' ;--- (S. Saar) 
progver = 'v2017-Mar-01' ;--- (J. Slavin)

prognam='XRT_SPOTCOR'
;t11=systime(1)

index_out = index
q_qt=keyword_set(quiet)
q_abort=keyword_set(abort)
q_all=keyword_set(all)
q_dust=n_elements(dust) ne 0  ; does dust array exist?
abort=0                       ; set abort flag

hist_val = strarr(n_elements(index)) + $ 
     'No spot correction necessary.'         ; initialize history note

fw1=index.ec_fw1
fw2=index.ec_fw2
fw1t=index.ec_fw1_
fw2t=index.ec_fw2_
if fw2t eq 'Gband' or index.ec_imty_ eq 'dark' or (fw1+fw2) eq 0 then begin
   if q_qt ne 1 then print,'Dark or Gband ... no spot correction needed.'
   data_out = data
   spotmap = data *0
   numspots = 0
   return
endif         

path0=getenv('SSW_XRT')+'/idl/util/'

if q_dust eq 0 then begin
   dfile='dust_pix1.dat'             ; read in dust file
   dlen=1035
   dust =intarr(2048,2048)  
   openr,1,path0+dfile                
   arr=intarr(2,dlen)                  
   readf,1,arr
   close,1
   xx=reform(arr[0,*])                ; read in x & y coords of features
   yy=reform(arr[1,*])
   dust[xx,yy]=1                       ; construct integer full-frame image 
endif

restgenx, file=path0+'xrt_spotmasters', spots_in, $
   stime_st, stime_en, sver, ssrc
nep=(size(spots_in))[3]    ; get # spot epochs
   

;  get data dimensions, FOV, binning, & filters used
if (size(data))[0] eq 3 then N=(size(data))[3]  else N=1


;  check for a "valid" data set.  All the same binning, size, FOV 
bin=index.chip_sum
ubin=uniq(bin)
n1=index.naxis1
un1=uniq(n1)
n2=index.naxis2
un2=uniq(n2)
pr=index.rpos_row
unr=uniq(pr)
pc=index.rpos_col
unc=uniq(pc)
; 
sumpar=n_elements(ubin)+n_elements(un1)+n_elements(un2)+n_elements(unr)+ $
   n_elements(unc)
if sumpar gt 5 then begin
   if q_qt eq 0 then print,'NAUGHTINESS DETECTED! Your images are not all ' + $
     'the same SIZE, BINNING, and CCD POSITION. I will pout and return ....'
   abort=1
   return 
endif
n1=n1[0] 
n2=n2[0]
bin=bin[0]
y00=index[0].p1row
x00=index[0].p1col

fcal=0                                     ; cal flag
if n1 gt 2048 then begin                   ; cal image!
  datac = data                             ; save original
  fcal = 1                                 ; set cal flag
  data = data[64:64+2047,0:2047]           ; extract active area
  n1 = 2048                                ; reset size
  n2 = 2048
endif 


data_out = data

date = index.date_obs
tmn = min(anytim(date))
tmx = max(anytim(date))

xrt_dustgrow,index,dustg,dust,hist=histd_val   ; estimate dust growth (if any)

filt=intarr(N)-2
ifix=where((fw1 le 3) and (fw2 le 3) and (fw1+fw2 ne 0) and $
   (tmn ge anytim(stime_st[0])) and (index.ec_imtyp eq 0),nfix) ; need to fix?
dlim=[.02,.02,.02,.02,.02]   ; spot depth limit: correct the spot 
                                 ; if (avg(spot)/median(boundary)-1) >dlim
limdb=[0.1+fltarr(6)]    ; spot flatness limit: spline correct the spot 
                               ; if (max(edge)-min(edge))/median >limdbk 
rx=2048.                      ;  CCD maximum size   (cal frames?) 
ry=2048.                      ;  CCD maximum size   (cal frames?)
if n_elements(thresh) eq 0 then thresh=0.5 ; default threshold=0.5
if nfix gt 0 then begin         ; if there are contam spots
   epcount=0
   for j=0,nep-1 do begin                           ;  determine spot epoch
     if (tmn ge anytim(stime_st[j])) and (tmx lt anytim(stime_en[j])) then begin
         iep=j                                      ; index for epoch
         epcount=epcount+1                          ; count epochs spanned
      endif
   endfor
   if epcount ne 1 then begin                      ; check for multiple epochs
     if q_qt eq 0 then print,'NAUGTHINESS DETECTED! Images span >1 spot epochs.'
      abort=1
      return
   endif 
   i0=where(fw2t eq 'Al_mesh',ni0)
   if ni0 gt 0 then filt[i0]=0
   i1=where(fw1t eq 'Al_poly' and fw2t eq 'Open',ni1)
   if ni1 gt 0 then filt[i1]=1
   i2=where(fw2t eq 'Ti_poly',ni2)  ; also for Ti-Poly/Al-Poly
   if ni2 gt 0 then filt[i2]=2
   i3=where(fw1t eq 'C_poly',ni3)
   if ni3 gt 0 then filt[i3]=3
   i4=where(fw1t eq 'Be_thin',ni4)
   if ni4 gt 0 then filt[i4]=4
;   i5=where(fw2t eq 'Gband',ni5)   
;   if ni5 gt 0 then filt(i5)=-1
   i6=where(fw1t eq 'C_poly' and fw2t eq 'Ti_poly',ni6)
   if ni6 gt 0 then filt[i6]=4   ; if C-Poly/Ti-Poly -use Be-thin  (filt = 4)


;  determine unique filters 
;  if WL = Gband (filt = -1) - use *all* spots
   filt=filt(where(filt ne -2))
   isf=sort(filt)
   ufilt=filt[isf[uniq(filt[isf])]]
   nfilt=n_elements(ufilt)
endif else begin                        ; otherwise setup blank spot array
   nfilt=1
   spots_in=fltarr(rx,ry)
   iep=0
endelse

numspots = fltarr(8,nfilt)

;   if bin eq 4 then thresh = 0.20             ; new 4x4 threshold = 3 of 16
spotmap = intarr(n1,n2,nfilt)         ; setup output spotmap

for i=0,nfilt-1 do begin             ; loop over filters to fix
   if nfix eq 0 then begin
      k = 0
      filt = filt*0
   endif else k = ifix[i]  
   filtk = filt[k]                    ; get filter # for filter k
;   gb=filtk eq -1                   ; gband flag
   iflk = where(filt eq filtk,nflk)    ; find when filtk is used
      
   limk = dlim[k]                     ; get spot depth limit for filter k
   limdbk = limdb[k]                  ; get spot flatness limit for filter k
   spots0 = fltarr(rx+2,ry+2)  ; setup 2048^2 spots array padded @ edge 
                               ; w/zeros (zeros absorb dilation problems @edge)
   spots0[1,1] = reform(spots_in[*,*,iep])   ; insert spots 

   dielem = fltarr(3,3) + 1         ; setup dilation kernel 
   spots = dilate(spots0,dielem)  ; expand spots by 1 pix in all directions, 
   spots = spots[1:rx,1:ry] + dustg ; remove pad, add dust, dilated (if needed)
   spots = spots < 1            ; elim double counting
   numspots[0,k] = [n_elements(xx),total(spots0),total(spots)] ; add spot stats

   spots = spots[x00:x00+n1*bin-1,y00:y00+n2*bin-1] ; get FOV(unbinned size)

   spotb = rebin(spots*1.,n1,n2)   ; rebin to appropriate size (n1xn2)
   spots = spotb ge thresh         ; reset spots > threshold (default=0)
                                    ; (may want to reset to eg, 1/64. or similar
                                    ; so that small spot fractions in heavily
                                    ; binned data do not get corrected)
   ind = where(spots eq 1.,ns)   ; find spot pixels (expanded) in n1xn2 map 
;   print,' (expanded) total (n1xn2): ',total(spots)
   if ns gt 0 then begin
      numspots[3,k] = ns     ; n_spotpix in FOV after expansion & thresholding
      xy = array_indices(spots,ind)   ; find their x & y coords
      xx = reform(xy[0,*])
      yy = reform(xy[1,*])
      spot_sw = intarr(n1+2,n2+2)      ; make spot_sw, size spots +1 in all dir
      spot_sw[1,1] = spots                       ; drop in spots
      mapspot = label_region(spot_sw eq 1,/all)  ; generate map numbering spots
      mapspot = mapspot[1:n1,1:n2]               ; trim (this gets border spots)
      ispots = 0L                   
      h = histogram(mapspot,rev=ispots) ; find loc. of unique spots (ispots)
      nspot = n_elements(h) -1  ; number of spots (not incl. zero)
      numspots[4,k] = nspot     ; n_spots in FOV after expansion & thresholding
      x0 = findgen(n1)#replicate(1.,n2)  ; setup x & y coordinate maps
      y0 = replicate(1.,n1)#findgen(n2)
      nthick = 1                 ;  # of pixels thick outer bnd to use in calc.
      dielem = intarr(2*nthick+1,2*nthick+1)+1  ; setup dilation kernel 
      spots0 = fltarr(n1+nthick*2,n2+nthick*2)     ; pad spot array 
      spots0[nthick,nthick] = spots     
      sbnd = dilate(spots0,dielem)              
      ; dilate, and zero spots to get boundaries (= sbnd)
      sbnd = sbnd[nthick:nthick+n1-1,nthick:nthick+n2-1]*(spots eq 0)
                                                                    
      if bin eq 1 then begin
         sbin=erode(spots0,dielem)   ; erode spots to get original 
                                  ; (unexpanded) spots (better for 1x1 data)
         numspots[5,k]=total(sbin)     ; n_spotpix in FOV in original spots
      endif
      spotk = spots                       ; copy spot array in this filter
      for j=1L,nspot do begin             ; loop over spots in this filter 
         ispotj = ispots[ispots[j]:ispots[j+1]-1]  ; find pixels in spot j
         xjlim = minmax(x0[ispotj]) +[-2,2] ; extend box around spots by +/-2
         yjlim = minmax(y0[ispotj]) +[-2,2] ; and get xy coords of bounding box
         xjl = xjlim[0]>0
         xjh = xjlim[1] <(n1-1)
         yjl = yjlim[0]>0
         yjh = yjlim[1] <(n2-1)
         xw = xjh-xjl+1                     ; get spot window size (xw*yw)
         yw = yjh-yjl+1
         spotj = spots[xjl:xjh,yjl:yjh]     ; extract window around spot
         sbndj = sbnd[xjl:xjh,yjl:yjh]      ; also get outer boundary
         x0j = x0[xjl:xjh,yjl:yjh]
         y0j = y0[xjl:xjh,yjl:yjh]
         ibad = where(spotj eq 1,nsj)       ; final spot pixels in window
         igood = where(spotj eq 0)          ; final non-spot pixels in window
         ibnd = where(sbndj eq 1,nsbnd)     ; final boundary pixels in window
         if bin eq 1 then begin             ; if 1x1 data , also
            ; get original (unexpanded) spot
            sbinj = sbin[xjl+1:xjh+1,yjl+1:yjh+1] 
            isbin = where(sbinj eq 1,nsbin)    ; and its location
         endif else nsbin=0
         xgood = x0j[igood]                 ; xy coords of non-spot pixels
         ygood = y0j[igood]
         m = 0                     
         nuncorr = 0L               ; initialize counter of uncorrected spots
         npixuncorr = 0L            ; initialize counter of uncorrected pixels
         repeat begin           ; fix this spot in all images in this filter 
            wind0 = reform(data[xjl:xjh,yjl:yjh,iflk[m]]) ; extract window
            ; Fix for problem of bad (missing frame) pixels being included in
            ; interpolation region around spots JDS 2/28/2017
            mfr = where(wind0 eq -999., nmfr, comp=okfr, ncomp=nokfr)
            if nmfr gt 0 and nokfr gt 0 then begin
                ; if possible replace with median of only non-spot pixels
                irg = where((wind0 ne -999.) and (spotj eq 0), nrg)
                if nrg gt 0 then md = median(wind0[irg]) $
                    else md = median(wind0[okfr])
                wind0[mfr] = md
            endif
            wind = median(wind0,3)  ; median to remove spikes @spot boundary
            if m eq 0 then begin  ; if this is 1st image in this filter...
               spota = avg(wind[ibad])  ;  define avg data value in spot
               bndm = median(wind[ibnd])   ; median data value at spot boundary 
               mmb = minmax(wind[ibnd])    ; minmax along boundary
               if abs(spota-bndm)/bndm gt limk then begin ; if fractional spot
                                                      ; depth > limk...
                  if (mmb[1]-mmb[0])/bndm gt limdbk then $  ; set flag on 
                     fixflag=3 else fixflag=1        ; unevenness of boundary 
                  if nsj gt 30 and bin eq 1 and filtk ge 0 then $
                     fixflag=(fixflag+1)<3 
                            ; if spot area>30 and 1x1 data, flag=flag+1
               endif else begin
                  fixflag=0                ; if spot depth < limk set flag=0
                  spotk[ispotj]=0              ; zero out spot in this filter
               end
            endif 
            if q_all eq 1 then fixflag=3
            case 1 of 
               (fixflag eq 1): begin       ; if spot deep, boundary even:  
                  bndm=median(wind[ibnd])  ; replace spot with smoothed
                  wind0[ibad]=bndm         ; boundary median 
               end
               (fixflag eq 2): begin    ; if spot deep, bnd even + area>30
                  bndm=median(wind[ibnd])  ; & data are 1x1 & not WL: replace
                  wind0[isbin]=bndm  ; non-expanded spot w/smoothd bnd median
               end
               (fixflag eq 3): begin       ; if spot deep, boundary uneven:
                  zgood=double(wind[igood]) ; do a thin-plate spline 
                  res=grid_tps(xgood,ygood,zgood,ngrid=[xw,yw],delta=[1,1])
                  avnew =avg(res[ibad])     ; avg fixed data
                  avold = avg(wind0[ibad])  ; avg old data
;                  if gb and avnew lt 1.05*avold then wind0[ibad]=res[ibad] 
                  if avnew gt 0.95*avold then wind0[ibad]=res[ibad]
                                       ; if fix ok, apply interpolation (others)
               end
               else: fixflag=0      ; if spot not deep, do not correct
            endcase
            ; Don't replace missing frame pixels with interpolated data
            if nmfr gt 0 then wind0[mfr] = -999.
            data_out[xjl:xjh,yjl:yjh,iflk[m]]=wind0 ;insert fixed spot window
            m=m+1
         end until (m gt (nflk-1)) or (fixflag eq 0) ; until last in filter 
                                        ; or decide spot shouldnt be corrected 
         nuncorr = nuncorr + (fixflag eq 0) ; add 1 if a spot is not corrected 
         npixuncorr =npixuncorr +(fixflag eq 0)*(nsj*(bin ne 1) $
            + nsbin*(bin eq 1))             ; add nsj or nsbin if  "  "  "  " 
      
      endfor   ; end loop over spots in this filter 
      if bin eq 1 then ns=total(sbin) ; if bin=1 , reset 
                                                   ;  # of spot pixels 
      numspots[6,k] = nspot - nuncorr    ; number of spots corrected in filter k
      numspots[7,k] = ns - npixuncorr   ; number of pixels  "   "   "   "
   endif
   spotmap[*,*,k]=spotb               ; original is output 
   if nfix gt 0 then begin
      if q_all eq 1 then histbeg='All spots and dust ' else $ 
         histbeg='Spots and dust ' 
      histend='Spot mastermap version = ' + sver[iep] + '. Threshold = ' + $
          string(thresh)
   endif else begin 
      histbeg='Dust '
      histend=''
   endelse
   hist_val[k] = histbeg + 'corrected with thin plate splines. ' + histend
endfor    ; end loop over filters

if fcal eq 1 then begin                              ; if cal frame, 
   data=datac                                        ; replace input
   datac[64,0] = data_out                            ; put 2048^2 result
   data_out = datac                                  ; into proper place
endif

 ;=== Now write the correction notes to the image HISTORY tag.
hist_prefix = prognam + ' ' + progver + ': (' + systim() + ') '
hist_entry = hist_prefix + hist_val + ' ' + histd_val
update_history, index_out, hist_entry, /noroutine, /mode

return
end
