% Experiment 3: Spatial domain filtering of image
% a. Contrast stretching  b. Histogram equalization
clear; close all; clc;

I = imread('pout.tif');                 % low-contrast image

%% Part A: Using inbuilt MATLAB functions (lab handout code)
J = histeq(I);                          % histogram equalization
K = imadjust(I);                        % contrast stretching

figure('Name', 'A_inbuilt');
subplot(2, 3, 1); imshow(I); title('Original Image');
subplot(2, 3, 2); imshow(K); title('Enhanced (imadjust)');
subplot(2, 3, 3); imshow(J); title('Histogram Equalized (histeq)');
subplot(2, 3, 4); imhist(I); title('Histogram: original');
subplot(2, 3, 5); imhist(K); title('Histogram: imadjust');
subplot(2, 3, 6); imhist(J); title('Histogram: histeq');

%% Part B: Custom implementation
[myK, rmin, rmax] = myContrastStretch(I, 0.01);
myJ = uint8(myHistEq(I, 256));
lims = stretchlim(I) * 255;
fprintf('Custom stretch limits : [%d %d]\n', rmin, rmax);
fprintf('stretchlim (imadjust) : [%.0f %.0f]\n', lims(1), lims(2));
fprintf('Mean |myContrastStretch - imadjust| = %.2f gray levels\n', mean(abs(double(myK(:)) - double(K(:)))));
fprintf('Std. deviation: original %.2f, stretched %.2f, equalized %.2f\n', ...
    std(double(I(:))), std(double(myK(:))), std(double(myJ(:))));

figure('Name', 'B_custom');
subplot(2, 3, 1); imshow(I);   title('Original Image');
subplot(2, 3, 2); imshow(myK); title('Custom contrast stretching');
subplot(2, 3, 3); imshow(myJ); title('Custom histogram equalization');
subplot(2, 3, 4); bar(0:255, myHistogram(I), 1, 'k');   xlim([0 255]); title('Histogram: original');
subplot(2, 3, 5); bar(0:255, myHistogram(myK), 1, 'k'); xlim([0 255]); title('Histogram: stretched');
subplot(2, 3, 6); bar(0:255, myHistogram(myJ), 1, 'k'); xlim([0 255]); title('Histogram: equalized');

% Transformation functions used by the two methods
r = 0:255;
sStretch = min(max((r - rmin) / (rmax - rmin) * 255, 0), 255);
sEq = myHistEq(uint8(r), 256, myHistogram(I));
figure('Name', 'B_transfer_functions');
plot(r, r, 'k:', r, sStretch, 'b-', r, sEq, 'r-', 'LineWidth', 1.5); grid on;
legend('Identity', 'Contrast stretching', 'Histogram equalization', 'Location', 'southeast');
xlabel('Input gray level r'); ylabel('Output gray level s'); title('Transformation functions');
axis([0 255 0 255]);

% Worked example from the handout: 3-bit image of size 4 x 4
exImg = [0 1 2 0; 3 4 7 2; 2 5 6 3; 0 3 4 0];
[exEq, exTable] = myHistEq(exImg, 8);
fprintf('\nWorked example (3-bit image, L = 8):\n');
fprintf('%5s %5s %9s %9s %5s\n', 'r_k', 'n_k', 'p_r', 'cdf', 's_k');
fprintf('%5d %5d %9.4f %9.4f %5d\n', table2array(exTable)');
fprintf('\nHistogram equalized 4 x 4 image:\n');
fprintf('%4d %4d %4d %4d\n', exEq');

%% Local functions
function h = myHistogram(img)
    % count the pixels at each gray level 0..255
    h = zeros(256, 1);
    img = double(img(:));
    for k = 1:numel(img)
        h(img(k) + 1) = h(img(k) + 1) + 1;
    end
end

function [out, rmin, rmax] = myContrastStretch(img, frac)
    % linear stretch between the gray levels that clip frac of pixels at each end
    c = cumsum(myHistogram(img)) / numel(img);
    rmin = find(c > frac, 1) - 1;
    rmax = find(c >= 1 - frac, 1) - 1;
    out = (double(img) - rmin) / (rmax - rmin) * 255;
    out = uint8(min(max(out, 0), 255));
end

function [out, T] = myHistEq(img, L, nk)
    % s_k = (L - 1) * sum_{j=0..k} p_r(r_j)
    img = double(img);
    if nargin < 3
        nk = zeros(L, 1);
        for k = 1:numel(img)
            nk(img(k) + 1) = nk(img(k) + 1) + 1;
        end
    end
    pr  = nk / sum(nk);
    cdf = cumsum(pr);
    sk  = round((L - 1) * cdf);
    out = reshape(sk(img + 1), size(img));
    T = table((0:L-1)', nk, pr, cdf, sk, 'VariableNames', {'r_k', 'n_k', 'p_r', 'cdf', 's_k'});
end