% Project 2
% Andy Doran
%
% function stockham(name,xdim,No,block)
%
% name  = 'input image' (as in 'lenna.256')
% xdim  = x dimension of input image (usually 256 or 512)
% No    = Variance of Noise
% block = Degraded image is broken up into block-by-block chunks
%         when computing Restoration Filter denominator
%
% This function takes an input image, runs it through a LSI filter h,
% and adds Gaussian noise to it.  The MSE between the degraded and
% original image is then calculated.  Power Spectrum Equalization
% is then performed and the degraded image is restored using the filter.
% The MSE between the restored and original image is then calculated and
% returned.  Suu in the PSE filter is estimated from the original
% image.  Svv in the PSE filter is estimated from the degraded image.
% h is never directly estimated, but we solve for the restoration
% filter G by breaking up the degraded image into blocks and using an
% approximation of the log of the FT of our blurring process.
% 

function stockham(name,xdim,No,block)

% Load image
pict = freadbin(name,xdim,xdim);

% Create LSI degradation model, need it to be phaseless
hi = 3.5^(-2);
h = zeros(256);
xl = 4;
xh = xdim - xl + 2;
h(1:xl,1:xl) = hi;
h(xh:xdim,1:xl) = hi;
h(1:xl,xh:xdim) = hi;
h(xh:xdim,xh:xdim) = hi;

% Plot blurring function spectra to make sure we have LPF
clf
subplot(221)
imagesc(fftshift(abs(fft2(h))));
colormap(gray);
txt = ['Blurring Function Spectra'];
title(txt)
axis square
axis off

% Create Gaussian noise, mean = 0, variance comes from input (No)
noise = sqrt(No)*randn(xdim,xdim);

% Run image through LSI Filter and then add noise
dpict = distimag(pict,h,noise);

% Plot degraded image
subplot(222)
imagesc(dpict);
colormap(gray);
txt = [num2str(name) ' with additive Gaussian Noise (mean=0, var=' , num2str(No), ')'];
title(txt)
axis square
axis off


% Calculate MSE of degraded image
error = dpict - pict;
sqerr = sum(sum(error.^2));
DMSE = sqerr/(xdim^2)

% Calculate power spectral density of input image for Numerator
DPICT = fft2(dpict);
Suu = abs(DPICT).^2;
NUM = sqrt(Suu);

% Calculate Denominator of Restoration Filter
DEN = ones(xdim);    % Initialize denominator
iter = xdim/block;          % Get number of iterations for loops
Nb = iter^2;                % Calculate total number of blocks

for i = 0:(iter-1)
  for j = 0:(iter-1)
    x1 = (i*block) + 1;
    x2 = (i+1)*block;
    y1 = (j*block) + 1;
    y2 = (j+1)*block;
    MINI = fft2(dpict(x1:x2,y1:y2),xdim,xdim);   % Zero Pad out to xdim
    MINI = abs(MINI);      % Get magnitude of Mini-block
    DEN = DEN.*(MINI.^(1/Nb));
  end
end

C = .57221;     % Euler's constant
const = exp((-C)/2);
DEN = const*DEN;

% Calculate Restoration Filter
G = NUM./DEN;

% Plot restoration spectra
subplot(224)
imagesc(abs(G));
colormap(gray);
txt = ['Restoration Filter Spectra'];
title(txt)
axis square
axis off

% Restore Image
DPICT = fft2(dpict);
RPICT = DPICT.*G;
rpict = ifft2(RPICT);

% Plot restored image
subplot(223)
imagesc(abs(rpict));
colormap(gray);
txt = [num2str(name) ' restored using Stockham Approach'];
title(txt)
axis square
axis off


% Calculate MSE of restored image
error = abs(rpict) - pict;
sqerr = sum(sum(error.^2));
RMSE = sqerr/(xdim^2)

