% Experiment 7: Image restoration - handling noise in image
% Arithmetic, Geometric, Harmonic, Contraharmonic mean filters,
% Median, Max, Min and Alpha-trimmed mean filters
clear; close all; clc;

f = im2double(imread('circuit.tif'));   % circuit board image (as in the lecture slides)
[M, N] = size(f);
rng(7);                                  % repeatable noise
m = 3; n = 3;                            % window size

% Noisy test images
gGauss  = imnoise(f, 'gaussian', 0, 0.01);
r = rand(M, N);
gPepper = f; gPepper(r < 0.1) = 0;       % pepper noise, probability 0.1
gSalt   = f; gSalt(r < 0.1) = 1;         % salt noise, probability 0.1
gSP     = imnoise(f, 'salt & pepper', 0.1);
gUni    = min(max(f + (rand(M, N) - 0.5) * 0.3, 0), 1);   % uniform noise
gUniSP  = imnoise(gUni, 'salt & pepper', 0.1);            % uniform + salt & pepper

%% Part A: Using inbuilt MATLAB functions
box = ones(m, n);
arithA = imfilter(gGauss, box / (m*n), 'replicate');
geoA   = exp(imfilter(log(gGauss + eps), box / (m*n), 'replicate'));
harmA  = (m*n) ./ imfilter(1 ./ (gSalt + eps), box, 'replicate');
Q = 1.5;
chPosA = imfilter(gPepper.^(Q+1), box, 'replicate') ./ (imfilter(gPepper.^Q, box, 'replicate') + eps);
chNegA = imfilter((gSalt + eps).^(-Q+1), box, 'replicate') ./ imfilter((gSalt + eps).^(-Q), box, 'replicate');
medA   = medfilt2(gSP, [m n], 'symmetric');
maxA   = ordfilt2(gPepper, m*n, box, 'symmetric');
minA   = ordfilt2(gSalt, 1, box, 'symmetric');
d = 6;                                   % alpha-trimmed: drop d/2 lowest and d/2 highest
alphaA = nlfilter(padarray(gUniSP, [2 2], 'symmetric'), [5 5], @(w) trimmean(w(:), 100 * d / 25));
alphaA = alphaA(3:end-2, 3:end-2);

figure('Name', 'A_mean_filters');
subplot(2, 4, 1); imshow(gGauss);  title('Gaussian noise');
subplot(2, 4, 2); imshow(arithA);  title('Arithmetic mean');
subplot(2, 4, 3); imshow(geoA);    title('Geometric mean');
subplot(2, 4, 4); imshow(gSalt);   title('Salt noise (0.1)');
subplot(2, 4, 5); imshow(harmA);   title('Harmonic mean (salt)');
subplot(2, 4, 6); imshow(chNegA);  title('Contraharmonic Q = -1.5 (salt)');
subplot(2, 4, 7); imshow(gPepper); title('Pepper noise (0.1)');
subplot(2, 4, 8); imshow(chPosA);  title('Contraharmonic Q = 1.5 (pepper)');

figure('Name', 'A_order_filters');
subplot(2, 4, 1); imshow(gSP);     title('Salt & pepper (0.1)');
subplot(2, 4, 2); imshow(medA);    title('Median (medfilt2)');
subplot(2, 4, 3); imshow(gPepper); title('Pepper noise');
subplot(2, 4, 4); imshow(maxA);    title('Max (ordfilt2)');
subplot(2, 4, 5); imshow(gSalt);   title('Salt noise');
subplot(2, 4, 6); imshow(minA);    title('Min (ordfilt2)');
subplot(2, 4, 7); imshow(gUniSP);  title('Uniform + salt & pepper');
subplot(2, 4, 8); imshow(alphaA);  title('Alpha-trimmed 5x5, d = 6');

%% Part B: Custom implementation
arithB = customMeanFilter(gGauss, m, n, 'arithmetic');
geoB   = customMeanFilter(gGauss, m, n, 'geometric');
harmB  = customMeanFilter(gSalt, m, n, 'harmonic');
chPosB = customMeanFilter(gPepper, m, n, 'contraharmonic', 1.5);
chNegB = customMeanFilter(gSalt, m, n, 'contraharmonic', -1.5);
medB   = customOrderFilter(gSP, m, n, 'median');
maxB   = customOrderFilter(gPepper, m, n, 'max');
minB   = customOrderFilter(gSalt, m, n, 'min');
alphaB = customOrderFilter(gUniSP, 5, 5, 'alpha', d);
arithUni = customMeanFilter(gUniSP, 5, 5, 'arithmetic');
medUni   = customOrderFilter(gUniSP, 5, 5, 'median');

figure('Name', 'B_custom');
subplot(3, 4, 1);  imshow(arithB);   title('Arithmetic (Gaussian)');
subplot(3, 4, 2);  imshow(geoB);     title('Geometric (Gaussian)');
subplot(3, 4, 3);  imshow(harmB);    title('Harmonic (salt)');
subplot(3, 4, 4);  imshow(chNegB);   title('Contraharmonic Q=-1.5 (salt)');
subplot(3, 4, 5);  imshow(chPosB);   title('Contraharmonic Q=1.5 (pepper)');
subplot(3, 4, 6);  imshow(medB);     title('Median (salt & pepper)');
subplot(3, 4, 7);  imshow(maxB);     title('Max (pepper)');
subplot(3, 4, 8);  imshow(minB);     title('Min (salt)');
subplot(3, 4, 9);  imshow(gUniSP);   title('Uniform + salt & pepper');
subplot(3, 4, 10); imshow(arithUni); title('Arithmetic 5x5');
subplot(3, 4, 11); imshow(medUni);   title('Median 5x5');
subplot(3, 4, 12); imshow(alphaB);   title('Alpha-trimmed 5x5, d = 6');

% PSNR table: inbuilt vs custom
names  = {'Arithmetic (Gaussian)', 'Geometric (Gaussian)', 'Harmonic (salt)', ...
          'Contraharmonic Q=1.5 (pepper)', 'Contraharmonic Q=-1.5 (salt)', ...
          'Median (salt & pepper)', 'Max (pepper)', 'Min (salt)', 'Alpha-trimmed (uniform+S&P)'};
noisy  = {gGauss, gGauss, gSalt, gPepper, gSalt, gSP, gPepper, gSalt, gUniSP};
inbuilt = {arithA, geoA, harmA, chPosA, chNegA, medA, maxA, minA, alphaA};
custom  = {arithB, geoB, harmB, chPosB, chNegB, medB, maxB, minB, alphaB};
fprintf('%-32s %8s %9s %8s\n', 'Filter', 'Noisy', 'Inbuilt', 'Custom');
for k = 1:numel(names)
    fprintf('%-32s %8.2f %9.2f %8.2f\n', names{k}, myPsnr(noisy{k}, f), myPsnr(inbuilt{k}, f), myPsnr(custom{k}, f));
end
fprintf('\nUniform + salt & pepper, 5x5 window: arithmetic %.2f dB, median %.2f dB, alpha-trimmed %.2f dB\n', ...
    myPsnr(arithUni, f), myPsnr(medUni, f), myPsnr(alphaB, f));

%% Local functions
function out = customMeanFilter(img, m, n, type, Q)
    P = padarray(img, [floor(m/2) floor(n/2)], 'replicate');
    out = zeros(size(img));
    for x = 1:size(img, 1)
        for y = 1:size(img, 2)
            w = reshape(P(x:x+m-1, y:y+n-1), [], 1);
            switch type
                case 'arithmetic'
                    out(x, y) = sum(w) / (m * n);
                case 'geometric'
                    out(x, y) = prod(w + eps)^(1 / (m * n));
                case 'harmonic'
                    out(x, y) = (m * n) / sum(1 ./ (w + eps));
                case 'contraharmonic'
                    w = w + eps;
                    out(x, y) = sum(w.^(Q + 1)) / sum(w.^Q);
            end
        end
    end
end

function out = customOrderFilter(img, m, n, type, d)
    P = padarray(img, [floor(m/2) floor(n/2)], 'symmetric');
    out = zeros(size(img));
    for x = 1:size(img, 1)
        for y = 1:size(img, 2)
            w = sort(reshape(P(x:x+m-1, y:y+n-1), [], 1));
            switch type
                case 'median'
                    out(x, y) = w(ceil(m * n / 2));
                case 'max'
                    out(x, y) = w(end);
                case 'min'
                    out(x, y) = w(1);
                case 'alpha'
                    out(x, y) = sum(w(d/2 + 1:end - d/2)) / (m * n - d);
            end
        end
    end
end

function p = myPsnr(a, ref)
    mse = mean((a(:) - ref(:)).^2);
    p = 10 * log10(1 / mse);
end