FORWARD_FUNCTION DIcal_denoise, SoftThreshold, BayesShrink
;;-----------------------------------------------------------------------------
;; PURPOSE:
;;	Used in the DI Calibration Pipeline.
;;	Removes random noise from the image by performing a soft thresholding
;;	in the wavelet domain. This is done in four steps.
;;
;;	1. Perform a wavelet transform
;;	2. Identify threshold T
;;	3. Perform soft thresholding by, 
;;		Wnew[i] = sign(W[i])(|W[i]|-T) or 0 (whichever is bigger)
;;	4. Perform reverse wavelet transform on Wnew
;;
;;	The threshold can be selected using the THRESH keyword, but if it is
;;	not, then the BayesShrink method is applied to determine the threshold
;;	for each subband (See BayesShrink function below). Generally, the
;;	BayesShrink method will be more effective (in the least squares sense)
;;	because different thresholds are applied to each subband. Whereas,
;;	using the THRESH keyword applied the same threshold to all subbands.
;;	
;;	However, BayesShrink requires some knowledge of the standard devation
;;	of the noise. If a priori knowledge of the noise is known, it is 
;;	recommended to use the SIGMA keyword. If SIGMA is not specified,
;;	the noise is estimated using the robust mean estimator described in:
;;	Johnstone, I and B. Silverman. 'Wavelet Threshold Estimators for Data 
;;		with Correlated Noise', Journal of the Royal Statistical 
;;		Society. Series B. Vol. 59, No. 2. (1997), pp. 319-351
;;	
;;
;; CALLING SEQUENCE:
;;	out = DIcal_denoise(in, THRESH=thresh, SIGMA=sigma)
;;
;; REQUIRED INPUTS:
;;	in - The image to which this pipeline element is going to be applied
;;	fitsHdr - The FITS header for the image
;;
;; OUTPUTS:
;;	RETURN - the image after going through this calibration step
;;
;; OPTIONAL INPUT KEYWORDS:
;;	THRESH - the threshold for wavelet coefficients to keep
;;	SIGMA - estimate of the standard deviation of the noise (default=1)
;;
;; EXAMPLE:
;;      IDL> imgOut = DIcal_denoise(imgIn)
;;
;; PROCEDURES USED (i.e. called directly!):
;;	SoftThreshold
;;	BayesShrink
;;
;; MODIFICATION HISTORY:
;;   2004-07-01  M. Desnoyer    Created
;;
;;-----------------------------------------------------------------------------

;;-----------------------------------------------------------------------------
;; PURPOSE:
;;	Performs a soft thresholding of the image where, for each pixel,
;;	pNew = max(sign(p)(|p|-thresh), 0)
;;
;; CALLING SEQUENCE:
;;	out = SoftThreshold(img, thresh)
;;
;; REQUIRED INPUTS:
;;	img - the image to threshold
;;	thresh - the threshold to apply
;;
;; OUTPUTS:
;;	RETURN - the image after being thresholded
;;
;; OPTIONAL INPUT KEYWORDS:
;;
;; EXAMPLE:
;;      IDL> imgOut = SoftThreshold(imgIn, thresh)
;;
;; PROCEDURES USED (i.e. called directly!)::
;;
;; MODIFICATION HISTORY:
;;   2004-07-07  M. Desnoyer    Created
;;   2005-04-12  M. Desnoyer    Reordered functions
;;
;;-----------------------------------------------------------------------------
FUNCTION SoftThreshold, img, thresh

sign = 2*(img GT 0)-1
out =  sign*((sign*img-thresh[0])>0)
;; Clear floating point errors
x = check_math(MASK=128)
RETURN, out

END

;;-----------------------------------------------------------------------------
;; PURPOSE:
;;	Performs soft thresholding on wavelet coefficients in all subbands.
;;	The threshold for each subband is determined using the BayesShrink
;;	method described in:
;;	Chang, G. et al., 'Adaptive Wavelet Thresholding for Image Denoising
;;		and Compression', IEEE Transactions on Image Processing.
;;		Vol. 9, No. 9. (Sept 2000). pp. 1532-1546
;;
;; CALLING SEQUENCE:
;;	out = BayesShrink(w, sigma)
;;
;; REQUIRED INPUTS:
;;	w - the wavelet coefficients to apply the BayesShrink to
;;	sigma - the standard deviation of the noise in the image
;;
;; OUTPUTS:
;;	RETURN - wavelet coefficients after being denoised
;;
;; OPTIONAL INPUT KEYWORDS:
;;
;; EXAMPLE:
;;      IDL> wOut = BayesShrink(wIn, sigma)
;;
;; PROCEDURES USED (i.e. called directly!)::
;;	SoftThreshold
;;
;; MODIFICATION HISTORY:
;;   2004-07-07  M. Desnoyer    Created
;;
;;-----------------------------------------------------------------------------
FUNCTION BayesShrink, w, sigma

dim = size(w, /dimensions)

;; Arrived at coarsest subband
IF size(w,/n_dimensions) LT 2 THEN RETURN, w

;; Identify the locations of the 4 subbands
half = dim/2
x = [[half[0],dim[0]-1],[half[0],dim[0]-1],[0,half[0]-1],[0,half[0]-1]]
y = [[half[1],dim[1]-1],[0,half[1]-1],[half[1],dim[1]-1],[0,half[1]-1]]

;; Perform thresholding on subbands HH, HL and LH
FOR i=0, 2 DO BEGIN
	sigmaY2 = total((w[x[0,i]:x[1,i],y[0,i]:y[1,i]])^2,/NAN)/half[0]/half[1]
	sigmaX = sqrt((sigmaY2-sigma^2)>0)
	thresh = (sigmaX EQ 0)? $
		max(abs(w[x[0,i]:x[1,i],y[0,i]:y[1,i]]),/NAN):sigma^2/sigmaX
	w[x[0,i]:x[1,i],y[0,i]:y[1,i]] = $
		SoftThreshold(w[x[0,i]:x[1,i],y[0,i]:y[1,i]], thresh)
ENDFOR

;; Do BayesShrink on LL
w[x[0,3]:x[1,3],y[0,3]:y[1,3]]= $
	BayesShrink(w[x[0,3]:x[1,3],y[0,3]:y[1,3]], sigma)

RETURN, w
END

;;-----------------------------------------------------------------------------
;; Main function
;;-----------------------------------------------------------------------------

FUNCTION DIcal_denoise, in, fitsHdr, THRESH=thresh, SIGMA=sigma


;; Do error handling 
;CATCH, error
error=0

IF error NE 0 THEN BEGIN
	CATCH, /CANCEL
	message, 'Denoise - ' + !ERROR_STATE.MSG, /NONAME
ENDIF

;; Get the value to insert for saturated pixels while doing calculcations
satVal = sxpar(fitsHdr, 'PSATVAL', COUNT=c1)

;; Make sure we got all pramaters from the FITS header
IF (c1 EQ 0) THEN message, 'Invalid FITS header', /NONAME

IF size(in, /n_dimensions) NE 2 THEN message, 'Invalid Image', /NONAME

;; Pad the image to make it's dimensions a power of two
inDim = size(in,/dimensions)
dim = 2^ceil(alog(inDim)/alog(2))
img = dblarr(dim[0],dim[1])
img[0:inDim[0]-1, 0:inDim[1]-1] = double(in)

;; Find all the saturated pixels and replace with satVal temporarily
satLoc = where(finite(img,/INFINITY) EQ 1,satCnt)
IF satCnt GT 0 THEN img[satLoc] = satVal

;; Replace all gaps with 0 temporarily
gapLoc = where(finite(img,/NAN) EQ 1, gapCnt)
IF gapCnt GT 0 THEN img[gapLoc] = 0

;; Perform wavelet transformation
w = wtn(img, 12, /double)

;; Threshold the wavelet coefficients
IF keyword_set(THRESH) EQ 0 THEN BEGIN
	;; Determine standard deviation of the noise
	if keyword_set(sigma) EQ 0 THEN BEGIN
		sig = median(abs(w[dim[0]/2:*, dim[1]/2:*]))/0.6745
	END ELSE sig=sigma
	;; Perform thresholding on each subband
	w = BayesShrink(w, sig)
ENDIF ELSE w = SoftThreshold(w, thresh)

;; Perform reverse wavelet transformation
img = wtn(w,12,/INVERSE)

;; Replace saturated pixels
IF satCnt GT 0 THEN img[satLoc] = !VALUES.D_INFINITY

;; Replace gaps
IF gapCnt GT 0 THEN img[gapLoc] = !VALUES.D_NAN

;; Update the FITS header
fxaddpar, fitsHdr, 'DENOISE', 'T'
fxaddpar, fitsHdr, 'DENOISEV', 'BAYESSHRINK'

RETURN, img[0:inDim[0]-1, 0:inDim[1]-1]

END