PRO restele_ima,in,out,pin,pout

; RESTORATION of an image (2D array) for
; the 2-D theoretical MTF of a diffraction-limited telescope.
; This version is used to find Wiener filter parameters and
; to check the effect of restoration on single frames.
; IN = input image (integer array, normalized to 10000)
; OUT = output image (integer array)
; (The arrays must have xdim, ydim = even number)
; Optional outputs: pin = 2D power spectrum of IN
;		    pout= 2D power spectrum of OUT
; 
; Note: program asks for Wiener filter parameters.
; In case of constant S/N, C=15 seems to be adequate.
; In case of variable S/N the approximation is:
; 	1/mtf(f)   for f le fo (in arcsec^-1)
;	mtf(f)/(mtf(f)^2+K*(f-fo))   for f gt fo
; i.e., increasing K we do more filtering.
;
; Routines called: NINT, MEAN, STDEV, SLIDE_IMAGE(IDL), FFT(IDL)
;
; Version of RESTELE.PRO, 26 Feb 2004, Michal


on_error,1

; PARAMETERS OF RESTORATION ---------------------------------------------

	stp=0.05448       ; SCALE (PIXEL SIZE) OF THE IMAGE  (arcsec/pix)
	lambda=4.305d-5   ; WORKING WAVELENGTH (cm)
	d=50.		  ; DIAMETER OF THE TELESCOPE (cm)

; image size --------------------------------------------------------

si=size(in)
if si(0) ne 2 then message,'Error - Input must be a 2D image'
xdim=si(1)
ydim=si(2)
print,'Image size (pixels):',xdim,ydim

n=xdim>ydim	; taking larger of the dimensions, so that
		; the basic working frame size will be n,n

a=findgen(n/2+1)/(n*stp)                   ; scale array for plotting

; 1-D calculation of MTF ------------------------------------------------

mtf=dblarr(n)
cc=lambda*2.062648d5/(d*n*stp)   ; 206264.8 arcsec in radian
ss=0

   for i=1,n do begin
	      
      omega=cc*(i-1)
      if omega lt 1.d0 then begin
        tt=acos(omega)-omega*sqrt(1.-omega*omega)
      endif else begin
        ss=ss+1
	if ss eq 1 then print,'MTF cutoff: ',i-1
	tt=0.d0
      endelse
      mtf(i-1)=2.*tt/3.141592654d0

   endfor

window,2,ys=640
plot,a,mtf,linestyle=5,xra=[-0.2,12],yra=[-0.2,4],/xst,/yst
print,''

; creating the 1-D Wiener filter ----------------------------------------

newc:
print,'Wiener filter with constant S/N (=0) or variable (=1) ?'
if get_kbrd(1) eq '0' then begin
	read,'Filter parameter C = ',c
	t1=(c+1)*mtf/(c*mtf*mtf+1)
endif else begin
	read,'Filter parameters K, fo : ',k,fo
	fp=nint(fo*n*stp)	; conversion to pixels
	t1=dblarr(n)		; filter
	uu=findgen(n-fp-1)+1.	; linearly increasing N/S
	uu=uu/n/stp
	t1(0:fp)=1./mtf(0:fp)
	t1(fp+1:n-1)=mtf(fp+1:n-1)/(mtf(fp+1:n-1)^2+k*uu)
endelse

wset,2
plot,a,t1,xra=[-0.2,12],yra=[-0.2,4],/xst,/yst,/noerase
print,'Change filter? (y/n)'
if get_kbrd(1) eq 'y' then goto,newc

print,'Do you want to stop here? (y/n)'
if get_kbrd(1) eq 'y' then begin
	out=0
	goto,fin
endif
read,'Width of apodized edges in % of image size: ',perc
print,'I am working...'

t1=float(t1)

; transition to 2-D Wiener filter t -------------------------------------

t=fltarr(n,n)
f1=fltarr(n/2+1,n/2+1)

cros0=min(where(t1 le 0))		; zero crossing point of t1,
					; i.e. we can compute only
for j=0,cros0 do begin			; the non-zero part. The rest
	for i=0,cros0 do begin		; are zeros by definition.
		ii=float(i)
		jj=float(j)
		r=sqrt(ii*ii+jj*jj)
		r1=fix(r)
		if r1 ge cros0 then begin
		  f1(i,j)=0.
		endif else begin
		  f1(i,j)=t1(r1)+(t1(r1+1)-t1(r1))*(r-r1)
		endelse
	endfor
endfor

f2=reverse(f1(1:n/2-1,*),1)

t(0:n/2,0:n/2)=f1			; t - 2-D Wiener filter
t(n/2+1:n-1,0:n/2)=f2
t(n/2+1:n-1,n/2+1:n-1)=reverse(f2(*,1:n/2-1),2)
t(0:n/2,n/2+1:n-1)=reverse(f1(*,1:n/2-1),2)
f1=0
f2=0
print,'   2-D Wiener filter created'

; creating apodization window apw ----------------------------------------

   npx=nint(xdim*perc/100.)
   npy=nint(ydim*perc/100.)

   ux=fltarr(xdim)+1.		; x-bell
   ux(0)=0.
   for i=1,npx do ux(i)=(1.-cos(!pi*i/npx))/2.
   ux(xdim-npx:xdim-1)=reverse(ux(1:npx))

   uy=fltarr(ydim)+1.		; y-bell
   uy(0)=0.
   for i=1,npy do uy(i)=(1.-cos(!pi*i/npy))/2.
   uy(ydim-npy:ydim-1)=reverse(uy(1:npy))

   apw=ux#uy		; apodization window (xdim,ydim)
   print,'   Apodization window created'


; RESTORATION ------------------------------------------------------------

   ima=in

; apodization

   mea=mean(ima)
   ima=ima-mea
   ima=ima*apw
   ima=ima+mea	 	; apodized frame
  
; creating working array (n,n)

   im=fltarr(n,n)+mea	; working array with mean(ima)
   im(0,0)=ima		; insert image

; 2-D Fourier filtering
   
   f1=fft(im,-1)

   ;power spectra:
   pre=float(f1)
   pim=imaginary(f1)
   pin=pre^2+pim^2
   pre=pre*t
   pim=pim*t
   pout=pre^2+pim^2
   window,6,xs=512,title='IDL 6: Power spectrum'
   contour,pin, /xst,/yst,lev=[0,0.025]
   wait,2
   contour,pout,/xst,/yst,lev=[0,0.025]

   out=float(fft(f1*t,1))   ; restored frame (real part)

   out=nint(out(0:xdim-1,0:ydim-1)) ; back to original size and type

std=stdev(in,mea)
print,'Original image contrast =',std/mea
std=stdev(out,mea)
print,'Restored image contrast =',std/mea

slide_image,bytscl(in, top=!d.n_colors),title='ORIGINAL', $
	show_full=0,xvis=512,yvis=512
slide_image,bytscl(out,top=!d.n_colors),title='RESTORED', $
	show_full=0,xvis=512,yvis=512
print,'Please, close scrolling windows manually.'

fin:
print,'End of program.'
END