FUNCTION xrt_fft_convolution, x, y, correlation=correlation, nopad=nopad

; =========================================================================
;+
; PROJECT:
;
;       Solar-B / XRT
;
; NAME:
;
;       XRT_FFT_CONVOLUTION
;
; CATEGORY:
;
;       Math utils
;
; PURPOSE:
;
;       Perform 1-dimensional convolution (optionally correlation) by using 
;       the Fast Fourier Transform (FFT) and appropriate padding to avoid 
;       edge effects.
;
; CALLING SEQUENCE:
;
;       images_convolved = XRT_FFT_CONVOLUTION(x, y, nopad=nopad)
;
;       images_correlated = XRT_FFT_CONVOLUTION(x, y, nopad=nopad, /correlation)
;
; INPUTS:
;
;       X - [Mandatory] (float or double precision 1D array, [Nx])
;           One of the two images to be convolved (or correlated).
;
;       Y - [Mandatory] (float or double precision 1D array, [Ny])
;           One of the two images to be convolved (or correlated).
;
; KEYWORDS:
;
;       /NOPAD       - [Optional input] (byte, integer, or long scalar)
;                      If this keyword is set, then no zero padding will be 
;                      performed on the arrays for the FFT.
;
;       /CORRELATION - [Optional input] (byte, integer, or long scalar)
;                      If this keyword is set, then the correlation of X and Y
;                      instead of their convolution will be calculated.
;
; OUTPUTS:
;
;       Return - [Mandatory] (float or double precision 1D array)
;                This will be the convolution of the inputs X and Y, unless
;                the /CORRELATION keyword is set, in which case the correlation
;                array of X and Y will be returned.
;
; COMMON BLOCKS:
;
;       none
;
; NOTES:
;
;       none
;
; CONTACT:
;
;       Comments, feedback, and bug reports regarding this routine may be
;       directed to this email address:
;                boerner ~at~ lmsal.com
;
; MODIFICATION HISTORY:
;
;progver = 'v2009-Mar-11' ;--- (P.Grigis (SAO)) Written.
;progver = 'v2012-Jan-03' ;--- (M.Weber (SAO)) Added a return line to the 
;                             nopad=1 conditional section.
;progver = 'v2012-Jan-13' ;--- (M. Cheung (LMSAL)) Corrected calls to pg_fft_2dim_convolution to aia_fft_2dim_convolution
;
progver = 'v2017-Apr-14' ;--- (P. Jibben (SAO)) Copied the AIA routines to the XRT directory and changed name for
;                              XRT use.
;
;-
; =========================================================================

  nx=n_elements(x)
  ny=n_elements(y)

  n=max([nx,ny])

  IF keyword_set(nopad) THEN BEGIN
    ;; This IF-THEN block is for the case where **no** zero-padding is indicated.

    IF keyword_set(correlation) THEN BEGIN
       fftres=fft(fft(x,-1)*conj(fft(y,-1)),1)*n
    ENDIF $
    ELSE BEGIN
       fftres=fft(fft(x,-1)*fft(y,-1),1)*n
    ENDELSE

    RETURN,(shift(fftres,n/2))

  ENDIF $
  ELSE BEGIN
    ;; This ELSE block is for the case where zero-padding is indicated.

    neff=2LL^(ceil(alog(1.5*n)/alog(2)))

    xpad=fltarr(neff)
    ypad=fltarr(neff)

    xpad[neff/2-nx/2]=x
    ypad[neff/2-ny/2]=y

    IF keyword_set(correlation) THEN BEGIN
       fftres=fft(fft(xpad,-1)*conj(fft(ypad,-1)),1)*neff
    ENDIF $
    ELSE BEGIN
       fftres=fft(fft(xpad,-1)*fft(ypad,-1),1)*neff
    ENDELSE

    RETURN,(shift(fftres,neff/2))[neff/2-n/2:neff/2+n/2-1]

  ENDELSE


END


; =========================================================================
; =========================================================================

FUNCTION xrt_fft_2dim_convolution, x, y, correlation=correlation, $
                                   nopad=nopad, xpad=xpad, ypad=ypad


; =========================================================================
;+
; PROJECT:
;
;       Solar-B / XRT
;
; NAME:
;
;       XRT_FFT_2DIM_CONVOLUTION
;
; CATEGORY:
;
;       Math utils
;
; PURPOSE:
;
;       Perform 2-dimensional convolution (optionally correlation) by using 
;       FFT and appropriate padding to avoid edge effects.
;
; CALLING SEQUENCE:
;
;       images_convolved = XRT_FFT_2DIM_CONVOLUTION(x, y, nopad=nopad)
;
;       images_correlated = XRT_FFT_2DIM_CONVOLUTION(x, y, nopad=nopad, /correlation)
;
; INPUTS:
;
;       X - [Mandatory] (float or double precision 1D array, [Nx])
;           One of the two images to be convolved (or correlated).
;
;       Y - [Mandatory] (float or double precision 1D array, [Ny])
;           One of the two images to be convolved (or correlated).
;
; KEYWORDS:
;
; KEYWORDS:
;
;       /CORRELATION - [Optional input] (byte, integer, or long scalar)
;                      If this keyword is set, then the correlation of X and Y
;                      instead of their convolution will be calculated.
;
;       /NOPAD       - [Optional input] (byte, integer, or long scalar)
;                      If this keyword is set, then no zero padding will be 
;                      performed on the arrays for the FFT.
;
;       XPAD         - [Optional output] (float or double 2D array)
;                      This is the padding array for input X.
;
;       YPAD         - [Optional output] (float or double 2D array)
;                      This is the padding array for input Y.
;
;
; OUTPUTS:
;
;       Return - [Mandatory] (float or double precision 2D array)
;                This will be the convolution of the inputs X and Y, unless
;                the /CORRELATION keyword is set, in which case the correlation
;                array of X and Y will be returned.
;
; COMMON BLOCKS:
;
;       none
;
; NOTES:
;
;       none
;
; CONTACT:
;
;       Comments, feedback, and bug reports regarding this routine may be
;       directed to this email address:
;                boerner ~at~ lmsal.com
;
; MODIFICATION HISTORY:
;
;progver = 'v2009-Mar-11' ;--- (P.Grigis (SAO)) Written.
;progver = 'v2009-Sep-28' ;--- (P.Grigis (SAO)) Investigate padding to see if
;                             it's correct...
;
progver = 'v2017-Apr-14' ;--- (P. Jibben (SAO)) Copied the AIA routines to the XRT directory and changed name for
;                              XRT use.
;-
; =========================================================================

  ;; Find sizes of input arrays.
  sx=size(x)
  sy=size(y)

  nx1=sx[1]
  ny1=sx[2]
  nx2=sy[1]
  ny2=sy[2]

  n1=max([nx1,nx2])
  n2=max([ny1,ny2])


  IF keyword_set(nopad) THEN BEGIN
     ;; This IF-THEN block is for the case where **no** zero-padding is indicated.

     IF keyword_set(correlation) THEN BEGIN
        fftres=fft(fft(x,-1)*conj(fft(y,-1)),1)*n1*n2
     ENDIF $
     ELSE BEGIN
        fftres=fft(fft(x,-1)*fft(y,-1),1)*n1*n2
     ENDELSE

     RETURN,(shift(fftres,n1/2,n2/2))

  ENDIF $
  ELSE BEGIN
     ;; This ELSE block is for the case where zero-padding is indicated.

     neff1=2LL^(ceil(alog(1.5*n1)/alog(2)))
     neff2=2LL^(ceil(alog(1.5*n2)/alog(2)))

     print,neff1,neff2

     xpad=fltarr(neff1,neff2)
     ypad=fltarr(neff1,neff2)

     xpad[neff1/2-nx1/2,neff2/2-ny1/2]=x
     ypad[neff1/2-nx2/2,neff2/2-ny2/2]=y

     IF keyword_set(correlation) THEN BEGIN
        fftres=fft(fft(xpad,-1)*conj(fft(ypad,-1)),1)*neff1*neff2
     ENDIF $
     ELSE BEGIN
        fftres=fft(fft(xpad,-1)*fft(ypad,-1),1)*neff1*neff2
     ENDELSE

     RETURN,(shift(fftres,neff1/2,neff2/2))[neff1/2-n1/2:neff1/2+n1/2-1,neff2/2-n2/2:neff2/2+n2/2-1]

  ENDELSE


END


; =========================================================================
; =========================================================================


FUNCTION xrt_deconvolve_richardsonlucy, image, psf, niter=niter, $
             error=error, verbose=verbose, ihat=ihat, ohat=ohat, $
             inputohat=inputohat

; =========================================================================
;+
; PROJECT:
;
;       Solar-B / XRT 
;
; NAME:
;       
;       XRT_DECONVOLVE_RICHARDSONLUCY
;
; CATEGORY:
;       
;       Image processing
;
; PURPOSE:
;       
;       Perform the Richardson-Lucy (RL) deconvolution of an image, given a 
;       Point Spread Function (PSF). (This routine should only be used in the
;       case where the PSF is known and constant across the image. In cases 
;       where the PSF is not known, a different, blind deconvolution scheme 
;       should be used instead.)
;
; CALLING SEQUENCE:
;
;       deconvolved_image = xrt_deconvolve_richardsonlucy(image, psf, niter=niter)
;       
; INPUTS:
;       
;       IMAGE - [Mandatory] (double or float array, [Ncols,Nrows])
;               The image or spectrum to be deconvolved.
;
;       PSF   - [Mandatory] (double or float array, [Ncols,Nrows])
;               An array with the same number of dimensions as IMAGE,
;               representing the Point Spread Function of the system that 
;               produced the image. PSF can have a different number of 
;               elements than IMAGE.
;
; KEYWORDS:
;       
;
;       NITER     - [Optional input] (integer or long scalar)
;                   The number of iterations to be performed (default: 25).
;
;       /VERBOSE  - [Optional input] (Boolean scalar) 
;                   If set, a record of the workings of the routine is
;                   printed to the screen.
;
;       INPUTOHAT - [Optional input] (double or float complex array, [Ncols,Nrows])
;                   Allows the user to input the OHAT array. This can be used,
;                   e.g., to restart the iteration if not enough runs were 
;                   performed.
;
;       ERROR     - [Optional output] (integer scalar)
;                   The error state from the routine. 
;                   The meaning of the values is described below:
;                     0: successful completion
;                     1: IMAGE or PSF have unallowed size 
;                        (0 or more than 2 dimensions)
;                     2: IMAGE or PSF have unallowed type 
;                        (floating point or double required)
;
;       IHAT      - [Optional output] (double or float complex array, [Ncols,Nrows])
;                   The final blurred object model.
;
;       OHAT      - [Optional output] (double or float complex array, [Ncols,Nrows])
;                   The final object model.
;
; OUTPUTS:
;       
;       Return - [Mandatory] (double or float array, [Ncols,Nrows])
;                The deconvolved image, or -1 if an error has occurred.
;
; COMMON BLOCKS:
;
; 	none 
;
; NOTES:
;
;       1) The RL algorithm performed here follows closely the algorithm
;          explained in P. A. Jansson, Deconvolution of Images and Spectra, 
;          2nd edition, 1997, Chapter 10.
;
;       2) This routine should only be used in the case where the Point Spread
;          Function (PSF) is known and constant across the image. In cases 
;          where the PSF is not known, a different, blind deconvolution scheme 
;          should be used instead.
;
;       3) This program does **not** perform any adjustment of the
;          PSF for image binning. The user is responsible for
;          correctly preparing a binned PSF input if one is required
;          for the image.
;
; CONTACT:
;
;       Comments, feedback, and bug reports regarding this routine may be 
;       directed to this email address: 
;                boerner ~at~ lmsal.com
;       
; MODIFICATION HISTORY:
;
;progver = 'v2009-Mar-11' ;--- (P.Grigis (SAO)) Written. 
;progver = 'v2009-Apr-16' ;--- (P.Grigis (SAO)) Allowed complex ohat input/output.
;progver = 'v2009-Sep-28' ;--- (P.Grigis (SAO)) Note that padding seems to be
;                             broken - always use /nopad until fixed.
;progver = 'v2011-Jul-13' ;--- (P.Grigis (SAO)) Made /nopad the default.
;
progver = 'v2017-Apr-14' ;--- (P. Jibben (SAO)) Copied the AIA routines to the XRT directory and changed name for
;                              XRT use.
;                      
;-
; =========================================================================


  ;; Error status starts at "no error".
  error=0

  ;; Padding functionality doesn't work, so turn off.
  nopad=1

  ;; Set default number of iterations.
  niter=fcheck(niter,5)

  ;; Get image size and dimension.
  imsize=size(image)
  psfsize=size(psf)

  ;; Check for scalars or 3D arrays that may sneak in (we don't want those).
  IF imsize[0] EQ 0 OR imsize[0] GT 2 OR psfsize[0] EQ 0 OR psfsize[0] GT 2 OR imsize[0] NE psfsize[0] THEN BEGIN
     print,'Invalid image or PSF size. This function works only with 1- or 2-dimensional'
     print,'arrays, and both arguments must have the same dimension.'
     error=1
     RETURN,-1
  ENDIF


  ;; Check for numeric floating point types. The test below is true if either 
  ;; IMAGE or PSF is not of type 4 or 5 (floating or double).
  IF ~(imsize[n_elements(imsize)-2]   EQ 4 || imsize[n_elements(imsize)-2]   EQ 5) || $
     ~(psfsize[n_elements(psfsize)-2] EQ 4 || psfsize[n_elements(psfsize)-2] EQ 5) THEN BEGIN
     print,'Invalid image or PSF type. This function works only with floating point'
     print,'or double numeric arrays.'
     error=2
     RETURN,-1
  ENDIF

  twodim=imsize[0] EQ 2

  ;; Print some general info on the screen.
  IF keyword_set(verbose) THEN BEGIN

     IF twodim THEN  BEGIN
        print,' '
        print,'Richardson-Lucy deconvolution (ver:'+progver+') of a '+strtrim(imsize[1],2)+'x'+strtrim(imsize[2],2)+' image' + $
              ' and a '+strtrim(psfsize[1],2)+'x'+strtrim(psfsize[2],2)+ ' PSF.'
        print,' '
     ENDIF $
     ELSE BEGIN
        print,' '
        print,'Richardson-Lucy deconvolution (ver:'+progver+') of a '+strtrim(n_elements(image),2)+ $
              '-element spectrum with a '+strtrim(n_elements(psf),2),'-element PSF'
        print,' '
     ENDELSE

  ENDIF


  ;; Perform main RL iteration.
  IF twodim THEN BEGIN
     ;; This section performs RL for 2-dim images.

     ;; PSF normalization.
     psfnorm=xrt_fft_2dim_convolution(psf,psf*0+1,nopad=nopad)

     ;; Starting values for RL iteration.
     IF (NOT keyword_set(inputohat)) OR n_elements(ohat) EQ 0 THEN BEGIN
        ;print,'NO Input OHAT!'
        ohat=image
     ENDIF


     FOR i=0,niter-1 DO BEGIN
        IF keyword_set(verbose) THEN $
           print,'Performing 2D iteration '+strtrim(i+1,2)+' of '+strtrim(niter,2)

        ;; These are the 2 simple steps of the RL algorithm.
        ihat=xrt_fft_2dim_convolution(psf,ohat,nopad=nopad)
        ohat*=xrt_fft_2dim_convolution(image/ihat,psf,/corr,nopad=nopad)/psfnorm
     ENDFOR


  ENDIF $
  ELSE BEGIN
     ;; This section performs RL for 1-dim images.

     ;; PSF normalization.
     psfnorm=xrt_fft_convolution(psf,psf*0+1,nopad=nopad)

     ;; Starting values for RL iteration.
     IF (~ keyword_set(inputohat)) OR n_elements(ohat) EQ 0 THEN ohat=image


     FOR i=0,niter-1 DO BEGIN
        IF keyword_set(verbose) THEN $
           print,'Performing 2D iteration '+strtrim(i+1,2)+' of '+strtrim(niter,2)

        ;; These are the 2 simple steps of the RL algorithm.
        ihat=xrt_fft_convolution(psf,ohat,nopad=nopad)
        ohat*=xrt_fft_convolution(image/ihat,psf,/corr,nopad=nopad)/psfnorm
     ENDFOR


  ENDELSE


  RETURN,abs(ohat)


END