function [ output_img_matrix ] = wiener( im, psf, varargin )
% Create a function that:
% ? can have 3 or 4 input parameters (use varargin):
%     ? if 3 input: input_img_matrix, psf, noise_to_signal_power_ratio
%     ? if 4 input: input_img_matrix, psf, noise_autocorr, signal_autocorr
%     ? output: Wiener-filtered (restored) image,
% ? it should work with grayscale images,
% ? output_img_matrix = | IFT( R ? FT( input_img_matrix ) ) |where
%     ? FT means Fourier Transform,
%     ? IFT means inverse Fourier Transform,
%     ? R stands for the Wiener-filter,
%     ? dot (?) in this case just a multiplication, since you are in frequency-domain,
%     ? |...| is just the absolute-value.

if (size(im,3) ~= 1)
    error 'This picture is not grayscale!'
end

if nargin == 3
    nsr = varargin{1};
elseif nargin == 4
    nsr = varargin{1} ./ varargin{2};
elseif isempty(varargin)
    error 'Not enough arguments'
else 
    error 'Too many arguments'
end

H = psf2otf(psf,size(im));
H_star = conj(H);
R = H_star ./ (H.^2 + nsr);

output_img_matrix = abs(ifft2(R .* fft2( im ) ) );

end