PRO saturation_fudge, image_in, sat, image_out, i_bleed, $
                      nsigm=nsigm

; =========================================================================
;+
; PROJECT:
;       Solar-B / XRT
;
; NAME:
;       SATURATION_FUDGE
;
; CATEGORY:
;       Data calibration
;
; PURPOSE:
;       Smooths over saturated areas by linearly connecting neighboring
;       unsaturated pixels and smoothing the result. This allows better 
;       removal of Fourier signals introduced into the data by the
;       read-out electronics (it reduces sinc ringing at sharp saturation 
;       edges).  Can optionally give locations of corrected saturation 
;       "bloom/bleed" pixels - pixels at the saturated region boundary 
;       which are affected by saturation but not saturated themselves.
;
; CALLING SEQUENCE:
;       SATURATION_FUDGE, image_in, sat, image_out, [i_bleed,]
;                     [nsigm=nsigm]
;
; INPUTS:
;       IMAGE_IN  - [Mandatory] (2-dim float array, [Nx,Ny])
;                   The image containing read-out signals.
;            SAT  - [Mandatory] can be either 
;                   a 2-dim byte array, [Nx,Ny] containing 1 if pixel 
;                   the pixel is saturated, or 0 otherwise 
;                   or
;                   scalar value of saturation level 
;
; KEYWORDS:
;         NSIGM   - [Optional] Number of standard deviations
;                   beyond which a pixel at the edge of the 
;                   saturated region must be above the local median  
;                   to considered bad  (affected by saturation 
;                   "bloom/bleed").
;                   Default is 0.2
;                   (Note: <=0 or much >1 not recommended)
;
; OUTPUTS:
;       IMAGE_OUT - [Mandatory] (2-dim float array, [Nx,Ny])
;                   The image with saturated pixels smoothed over.
;                   Optionally, if /stop_bleed is set, saturation
;                   "bloom/bleed" pixels are corrected.
;       I_BLEED   - [Optional] (1-dim long array, [Nbleed])
;                   Index of pixels where saturation "bloom/bleed"
;                   has been corrected.
;
; EXAMPLES:
;
;       Basic usage #1 (with input map of saturated pixel locations):
;       IDL> saturation_fudge, image_in, sat_map, image_out 
;
;       Basic usage #2 (with input saturation value (=2500)):
;       IDL> saturation_fudge, image_in, 2500, image_out 
;
;       You want the locations of corrected saturation "bloom/bleed" pixels:
;     IDL> saturation_fudge, image_in, sat_map, image_out, i_bleed 
;
; COMMON BLOCKS:
;       None.
;
; NOTES:
;       1) CCD saturation often causes a bleed of excess charge in
;       the vertical (y) direction on the chip.  If the bleed is
;       over areas with low DN, gradients can be large, causing
;       problems for the Fourier space removal of read-out patterns 
;       (sinc ringing in the output).  This procedure linearly connects
;       (in x) unsaturated pixels just outside the saturation border,
;       and smooths the resulting "patch" in x and y.
;
;       2) Since border pixels are often affected by charge bleed as well
;       (though not saturated themselves), these "bloom/bleed" pixels 
;       are located and replaced with a local median value. Optionally,
;       the locations of these corrected pixels (i_bleed) can be kept 
;       (e.g., for recording purposes, or to change them back to their
;       original values later).
;       
; MODIFICATION HISTORY:
;
progver = 'v2007.May.18' ; --- (SS, MW) Written by SS. Cleaned up by MW.
progver = 'v2013.Feb.01' ; --- (SS) Fixed overcorrection for cases with 
;                        ;       multiple boundaries in x.  
progver = 'v2016.Nov.22' ; --- (JS) Fixed bloom/bleed calculation to prevent 
;                                wraparound and ensure no double counts
progver = 'v2017.May.30' ; --- (JS) Fixed problem for cases where only one
                         ;     bloom/bleed pixel is found, which caused error
                         ;     in call to STDEV
;
;-
; =========================================================================


;===== define "bloom/bleed" level as >nsigm*sigma above local median 
if not keyword_set(nsigm) then nsigm=0.2  

siz=size(image_in)
nx=siz(1)
ny=siz(2)

;===== ID and number distinct regions with DN>sat_level (isat)
im2=image_in
nsat=n_elements(sat)
if nsat eq 1 then begin    ;=====  saturation level entered 
   satu= im2 ge sat
   isat=label_region(satu)
   sat_level=sat
endif else begin          ;=====  saturation array entered
   isat=label_region(sat)
   sat_level=max(sat)
   satu=sat
end

if total(satu) eq 0 then begin ;=====  return if no saturation present
  image_out = image_in
  i_bleed = -1L
  return
endif

;===== find possible saturation bleed/bloom (in y) pixels at boundaries
;=====   and replace with local median (with saturated pixels suppressed)
unsat = satu eq 0
mask_edge = satu*0 + 1
mask_edge[*, [0,ny-1]] = 0
bsat = mask_edge*(shift(satu,0,1)+shift(satu,0,-1))
bbnd = (bsat*unsat < 1)
ibbnd = where(bbnd eq 1,nbbnd)
im2m = median(im2*unsat,3)
if nbbnd gt 1 then begin
    ; if only 1 bb pixel then can't get a stdev
    ib = where(im2[ibbnd] ge (im2m[ibbnd] + nsigm*stdev(im2m[ibbnd])),nib)
    if nib gt 0 then begin
       ib = ibbnd[ib]
       im2[ib] = im2m[ib] + (im2m[ib] eq 0)*im2[ib]
    endif
endif else if nbbnd eq 1 then begin
    imm = im2m[ibbnd[0]]
    if im2[ibbnd[0]] > imm then im2[ibbnd[0]] = im2m[ibbnd[0]]
endif

;=====  get indicies in the various regions (revi)
nr=max(isat)
hr=histogram(isat,rev=revi)

;===== loop over regions, determine x,y indicies for points in each
for j=1,nr do begin
   irj=revi(revi(j):revi(j+1)-1)
   irxj= irj mod nx
   iryj= (irj - irxj)/nx 

;=====  step in y, connect unsaturated edges of saturated regions with lines 
   for k=min(iryj),max(iryj) do begin
      iryk=where(iryj eq k)
      xkmin=min(irxj(iryk)) 
      xkmax=max(irxj(iryk)) 
      xkmax0=xkmax
      xkbit=findgen(xkmax-xkmin+1)+xkmin
;===== look for end of next saturated section of region at row y=k 
;=====    set section endpoint, or if none, set at end  
      repeat begin
         iend=min(where(im2(xkmin:xkmax0,k) lt sat_level))
         iend=iend(0)
         if iend eq -1 then xkmax=xkmax0 else xkmax=xkbit(iend)-1

;===== if section hasnt yet been done, do the linear fits connecting 
;=====   saturation-unaffected endpoints
         if iend ne 0 then begin
            slope=(im2(xkmax+1,k)-im2(xkmin-1,k))/(xkmax+1-(xkmin-1.))
            icept=im2(xkmin-1,k)-slope*(xkmin-1)
            im2(xkmin:xkmax,k)=slope*xkbit(0:xkmax-xkmin) + icept
         endif 

;===== look for start of next saturated section of region at row y=k 
         if iend ne -1 then begin
            ibeg=min(where(im2(xkmax+1:xkmax0,k) ge sat_level))
            ibeg=ibeg(0) 
            if ibeg ne -1 then begin
               xkmin=xkmax+1+ibeg(0)
               xkbit=findgen(xkmax0-xkmin+1)+xkmin
            endif else xkmin=xkmax0+1
         endif else xkmin=xkmax0+1
      end until iend eq -1 or xkmin gt xkmax0
   endfor
endfor

;===== find where changes have been made (idiff), and boundary
;=====    "bloom/bleed" pixels (i_bleed)
;idiff=where(im2 ne image_in, count_idiff)
;i_bleed=where(im2*unsat ne image_in*unsat, count_i_bleed)
idiff=where(bsat eq 1, count_idiff)
i_bleed = ibbnd 
count_i_bleed=n_elements(i_bleed)

;===== smooth the linear saturation replacement "cap" and make output
im2s=smooth(median(im2,3),[1,5])
image_out= image_in
if (count_idiff ge 1) then image_out(idiff) = im2s(idiff)
if (count_i_bleed ge 1) then image_out(i_bleed) = im2(i_bleed)

return


END ;======================================================================