module ImageProc
	interface 
		subroutine importimage(filename, img)
			use kind_defs
			implicit none
			character(*), intent(in) :: filename
			real(dbl), intent(out) :: img(:,:)
		end subroutine
		
		subroutine saveimage(filename, img)
			use kind_defs
			implicit none
			character(*), intent(in) :: filename
			real(dbl), intent(in) :: img(:,:)
		end subroutine
		
		function fftshift(x)
			use kind_defs
			implicit none
			real(dbl), intent(in) :: x(:,:)
			real(dbl) :: fftshift(size(x,1),size(x,2))
		end function
		
		function roundfilter(N1, N2, filt)
			implicit none
			integer, intent(in) :: N1, N2, filt
			integer :: roundfilter(N1, N2)
		end function
		
		subroutine dftfilt(img, psf, filt)
			use kind_defs
			implicit none
			real(dbl), intent(inout) :: img(:,:)
			real(dbl), intent(in) :: psf(:,:)
			integer, intent(in) :: filt
		end subroutine
	end interface
end module

subroutine importimage(filename, img)
	use kind_defs
	implicit none
	character(*), intent(in) :: filename
	real(dbl), intent(out) :: img(:,:)

	open(50, file=trim(filename), status='old', action='read')
		read(50, *) img
	close(50)
end subroutine

subroutine saveimage(filename, img)
	use kind_defs
	implicit none
	character(*), intent(in) :: filename
	real(dbl), intent(in) :: img(:,:)
	integer :: i, N1, N2
	
	N1 = size(img,1)
	N2 = size(img,2)
	open(51, file=trim(filename), status='replace', action='write')
		do i = 1,N1
			write(51, *) img(i,1:N2)
		enddo
	close(51)
end subroutine

function fftshift(x)
	use kind_defs
	implicit none
	real(dbl), intent(in) :: x(:,:)
	real(dbl) :: fftshift(size(x,1),size(x,2))
	integer :: N1, N2

	! Get the dimensions of the image:
	N1 = size(x,1)
	N2 = size(x,2)
	
	! Shift the quadrants of the image:
	fftshift(1:N1/2, 1:N2/2) = x(N1/2+1:N1, N2/2+1:N2)
	fftshift(1:N1/2, N2/2+1:N2) = x(N1/2+1:N1, 1:N2/2)
	fftshift(N1/2+1:N1, 1:N2/2) = x(1:N1/2, N2/2+1:N2)
	fftshift(N1/2+1:N1, N2/2+1:N2) = x(1:N1/2, 1:N2/2)
end function

! Returns a filter to be used to multiply with the fourier transform
! before doing the inverse filtering.
! Input: N1   -- Integer -- Size of first dimension of Image
!        N2   -- Integer -- Size of second dimension of Image
!        filt -- Integer -- Size of the filter. 
function roundfilter(N1, N2, filt)
	implicit none
	integer, intent(in) :: N1, N2, filt
	integer :: roundfilter(N1, N2)
	integer :: i, j
	
	roundfilter = 0
	do i = 1,N1
		do j = 1,N2
		    if (sqrt(real(i**2+j**2)) <= filt .or. sqrt(real((i-N1)**2 + j**2)) <= filt &
		    	& .or. sqrt(real(i**2 + (j-N2)**2)) <= filt .or. sqrt(real((i-N1)**2 + (j-N2)**2)) <= filt) then
		        roundfilter(i,j) = 1
		    endif
		enddo
	enddo
end function


! Perform filtering for the Point Spread Function.
! Input:   img  -- inout -- The image to be filtered.
!          psf  -- in    -- The image of the PSF.
!          filt -- in    -- Integer-value of the size of the pseudoinverse-filter.
subroutine dftfilt(img, psf, filt)
	use kind_defs
	use FourierTransform
	use ImageProc, only : roundfilter !, saveimage
	implicit none
	real(dbl), intent(inout) :: img(:,:)
	real(dbl), intent(in) :: psf(:,:)
	integer, intent(in) :: filt
	real(dbl), allocatable :: f_img(:,:), f_psf(:,:)
	integer, allocatable :: filter(:,:)
	complex(dbl), allocatable :: c_img(:,:), c_psf(:,:)
	integer :: m, P, N1, N2, M1, M2

	! Get size of the image:
	N1 = size(img,1)
	N2 = size(img,2)
	M1 = size(psf,1)
	M2 = size(psf,2)

	! Make them the same size:
	m = maxval([N1, N2, M1, M2])
	P = 2**nextpow2(2*m)
	
	! Allocate the images:
	allocate(f_img(P,P), f_psf(P,P), c_img(P,P), c_psf(P,P), filter(P,P))
	
	! Pad with zeros:
	f_img = 0
	f_img(1:N1, 1:N2) = img
	f_psf = 0
	f_psf(1:M1, 1:M2) = psf

	! Calculate the FFT of the two images:
	print *, 'Calculation Fourier-transforms...'
	c_img = fft2(f_img)
	c_psf = fft2(f_psf)

	!print *, 'Saving images...'
	!call saveimage('c_img.img', abs(c_img))
	!call saveimage('c_psf.img', abs(c_psf))

	! Correct for the PSF:
	print *, 'Correcting for the PSF...'
	c_img = c_img / c_psf
	
	! To avoid large errors do to zeros in c_psf:
	! This is called pseudoinverse filtering.
	filter = roundfilter(P, P, filt)
	c_img = c_img*filter
	
	! Do the inverse transformation:
	f_img = real(ifft2(c_img))
	
	! Crop to the final image:
	img = f_img(1:N1, 1:N2)
	
	! Final adjustments to image:
	img = transpose(img)
	img = img - minval(img) ! Make the image positive
	
	! Clean up the memory:
	deallocate(f_img, f_psf, c_img, c_psf, filter)
end subroutine