% Experiment 8: Image restoration - handling noise in image
% a. Wiener filter  b. Adaptive mean (local noise reduction) filter  c. Adaptive median filter
clear; close all; clc;

%% Part A: Lab handout programs
% a. Image filtering using the Wiener filter (wiener2)
I = imread('cameraman.tif');
I_noisy = imnoise(I, 'gaussian', 0, 0.01);       % mean = 0, variance = 0.01
I_filtered = wiener2(I_noisy, [5 5]);           % [m n] neighbourhood size

figure('Name', 'A_wiener');
subplot(1, 3, 1); imshow(I);          title('Original Image');
subplot(1, 3, 2); imshow(I_noisy);    title('Noisy Image (Gaussian)');
subplot(1, 3, 3); imshow(I_filtered); title('Filtered Image (wiener2)');
wienerA = im2double(I_filtered);

% b. Adaptive mean filtering
I = im2double(imread('cameraman.tif'));
I_noisy = imnoise(I, 'gaussian', 0, 0.01);
gaussNoisy = I_noisy;
[m, n] = size(I_noisy);
I_filtered = zeros(m, n);
win_size = 5;                                   % 5 x 5 neighbourhood
pad_size = floor(win_size/2);
I_padded = padarray(I_noisy, [pad_size pad_size], 'symmetric');
noise_var = 0.01;                               % noise variance
for i = 1:m
    for j = 1:n
        local_window = I_padded(i:i+win_size-1, j:j+win_size-1);
        local_mean = mean(local_window(:));
        local_var = var(local_window(:));
        if local_var > 0
            I_filtered(i, j) = I_noisy(i, j) - (noise_var/local_var) * (I_noisy(i, j) - local_mean);
        else
            I_filtered(i, j) = local_mean;
        end
    end
end
adaptMeanA = I_filtered;

figure('Name', 'A_adaptive_mean');
subplot(1, 3, 1); imshow(I);          title('Original Image');
subplot(1, 3, 2); imshow(I_noisy);    title('Noisy Image');
subplot(1, 3, 3); imshow(I_filtered); title('Filtered Image (Adaptive Mean)');

% c. Adaptive median filter (function saved in adaptiveMedianFilter.m)
noiseDensity = 0.20;                            % 20% noise
I_noisy = imnoise(I, 'salt & pepper', noiseDensity);
spNoisy = I_noisy;
S_max = 7;                                      % maximum allowed window size
I_filtered = adaptiveMedianFilter(I_noisy, S_max);
adaptMedA = I_filtered;

figure('Name', 'A_adaptive_median');
subplot(1, 3, 1); imshow(I);          title('Original');
subplot(1, 3, 2); imshow(I_noisy);    title('Noisy (Salt & Pepper)');
subplot(1, 3, 3); imshow(I_filtered); title('Filtered (Adaptive Median)');

%% Part B: Custom implementation
% Wiener filter written from the local statistics (same model as wiener2)
[wienerB, estNoise] = myWiener2(gaussNoisy, 5);
% Adaptive local noise reduction with the ratio sigma_n^2 / sigma_L^2 limited to 1
adaptMeanB = myAdaptiveMean(gaussNoisy, 5, 0.01);
% Compare the adaptive median filter with a fixed 7 x 7 median filter
medFixed = medfilt2(spNoisy, [7 7], 'symmetric');
med3 = medfilt2(spNoisy, [3 3], 'symmetric');

fprintf('Estimated noise variance (custom Wiener) = %.4f (true value 0.01)\n', estNoise);
fprintf('\nGaussian noise, variance 0.01 (PSNR in dB)\n');
fprintf('%-38s %7.2f\n', 'Noisy image', psnr(gaussNoisy, I));
fprintf('%-38s %7.2f\n', 'wiener2 (inbuilt, 5x5)', psnr(wienerA, I));
fprintf('%-38s %7.2f\n', 'Custom Wiener (5x5)', psnr(wienerB, I));
fprintf('%-38s %7.2f\n', 'Adaptive mean (handout)', psnr(adaptMeanA, I));
fprintf('%-38s %7.2f\n', 'Adaptive mean (ratio limited to 1)', psnr(adaptMeanB, I));
fprintf('%-38s %7.2f\n', 'Arithmetic mean 5x5', psnr(imfilter(gaussNoisy, ones(5)/25, 'symmetric'), I));
fprintf('\nSalt & pepper noise, density 0.20 (PSNR in dB)\n');
fprintf('%-38s %7.2f\n', 'Noisy image', psnr(spNoisy, I));
fprintf('%-38s %7.2f\n', 'Median 3x3', psnr(med3, I));
fprintf('%-38s %7.2f\n', 'Median 7x7', psnr(medFixed, I));
fprintf('%-38s %7.2f\n', 'Adaptive median (S_max = 7)', psnr(adaptMedA, I));

figure('Name', 'B_custom');
subplot(2, 4, 1); imshow(gaussNoisy);  title('Gaussian noise (0.01)');
subplot(2, 4, 2); imshow(wienerA);     title('wiener2');
subplot(2, 4, 3); imshow(wienerB);     title('Custom Wiener');
subplot(2, 4, 4); imshow(adaptMeanB);  title('Custom adaptive mean');
subplot(2, 4, 5); imshow(spNoisy);     title('Salt & pepper (0.20)');
subplot(2, 4, 6); imshow(med3);        title('Median 3x3');
subplot(2, 4, 7); imshow(medFixed);    title('Median 7x7');
subplot(2, 4, 8); imshow(adaptMedA);   title('Adaptive median, S_{max} = 7');

%% Local functions
function [out, noiseVar] = myWiener2(g, w)
    % f = m_L + max(sigma_L^2 - sigma_n^2, 0) / sigma_L^2 * (g - m_L)
    pad = floor(w / 2);
    P = padarray(g, [pad pad], 'symmetric');
    [M, N] = size(g);
    mL = zeros(M, N); vL = zeros(M, N);
    for i = 1:M
        for j = 1:N
            win = P(i:i+w-1, j:j+w-1);
            mL(i, j) = mean(win(:));
            vL(i, j) = mean(win(:).^2) - mL(i, j)^2;
        end
    end
    noiseVar = mean(vL(:));                % noise estimated as the mean local variance
    out = mL + max(vL - noiseVar, 0) ./ max(vL, noiseVar) .* (g - mL);
end

function out = myAdaptiveMean(g, w, noiseVar)
    pad = floor(w / 2);
    P = padarray(g, [pad pad], 'symmetric');
    out = zeros(size(g));
    for i = 1:size(g, 1)
        for j = 1:size(g, 2)
            win = P(i:i+w-1, j:j+w-1);
            mL = mean(win(:));
            vL = var(win(:));
            ratio = min(noiseVar / max(vL, eps), 1);   % never exceeds 1
            out(i, j) = g(i, j) - ratio * (g(i, j) - mL);
        end
    end
end