% Experiment 9: Image compression using DCT, then calculation of PSNR and compression ratio
clear; close all; clc;

%% Part A: Lab handout program (global DCT, keep top X% coefficients)
% vision.DCT, vision.IDCT and vision.PSNR are no longer available in
% MATLAB R2026b, so dct2, idct2 and psnr are used in their place.
keepPercent = 40;                       % keep top 40% coefficients
I = im2double(imread('cameraman.tif'));
J = dct2(I);                            % global 2D DCT

% Find threshold for the top keepPercent coefficients
coeffs = abs(J(:));
sortedCoeffs = sort(coeffs, 'descend');
cutoffIndex = round(numel(sortedCoeffs) * keepPercent / 100);
threshold = sortedCoeffs(cutoffIndex);

% Zero out coefficients below threshold
J(abs(J) < threshold) = 0;

% Inverse DCT
It = idct2(J);

% PSNR and compression ratio
psnr_val = psnr(It, I);
fprintf('PSNR = %.2f dB\n', psnr_val);
nonzero = nnz(J);
total = numel(J);
CR = total / nonzero;
fprintf('Compression Ratio = %.2f\n', CR);

figure('Name', 'A_inbuilt');
subplot(1, 3, 1); imshow(I);  title('Original');
subplot(1, 3, 2); imshow(log(1 + abs(dct2(I))), []); colormap(gca, 'jet'); title('DCT coefficients (log)');
subplot(1, 3, 3); imshow(It); title(sprintf('Reconstructed (Top %d%%, PSNR=%.2f dB)', keepPercent, psnr_val));

imwrite(It, 'CameramanCompressed.tif');   % save compressed image

%% Part B: Custom implementation (own DCT matrix, PSNR and compression ratio)
N = size(I, 1);
T = myDctMatrix(N);
Jc = T * I * T';                          % 2D DCT, F = T f T'
fprintf('\nMax |custom DCT - dct2| = %.2e\n', max(abs(Jc(:) - reshape(dct2(I), [], 1))));

percents = [1 5 10 20 40 60];
recon = cell(1, numel(percents));
psnrs = zeros(1, numel(percents));
crs = zeros(1, numel(percents));
fprintf('\n%-10s %10s %12s %10s\n', 'Kept (%)', 'CR', 'PSNR (dB)', 'MSE');
for k = 1:numel(percents)
    Jk = Jc;
    s = sort(abs(Jk(:)), 'descend');
    th = s(round(numel(s) * percents(k) / 100));
    Jk(abs(Jk) < th) = 0;
    recon{k} = T' * Jk * T;               % inverse DCT, f = T' F T
    mse = sum((I(:) - recon{k}(:)).^2) / numel(I);
    psnrs(k) = 10 * log10(1 / mse);       % peak value is 1 for double images
    crs(k) = numel(Jk) / nnz(Jk);
    fprintf('%-10d %10.2f %12.2f %10.6f\n', percents(k), crs(k), psnrs(k), mse);
end

figure('Name', 'B_custom');
for k = 1:numel(percents)
    subplot(2, 3, k); imshow(min(max(recon{k}, 0), 1));
    title(sprintf('Top %d%%: CR = %.1f, PSNR = %.2f dB', percents(k), crs(k), psnrs(k)));
end

figure('Name', 'B_psnr_vs_cr');
plot(crs, psnrs, 'bo-', 'LineWidth', 1.5, 'MarkerFaceColor', 'b'); grid on;
for k = 1:numel(percents)
    text(crs(k), psnrs(k) + 1, sprintf('%d%%', percents(k)), 'HorizontalAlignment', 'center');
end
set(gca, 'XScale', 'log');
xlabel('Compression ratio (log scale)'); ylabel('PSNR (dB)');
title('Quality versus compression for the global DCT');

%% Local functions
function T = myDctMatrix(N)
    % orthonormal DCT-II basis: T(i,j) = a(i) cos(pi (2j+1) i / 2N)
    T = zeros(N, N);
    for i = 0:N-1
        for j = 0:N-1
            if i == 0
                T(i+1, j+1) = 1 / sqrt(N);
            else
                T(i+1, j+1) = sqrt(2/N) * cos((pi * (2*j + 1) * i) / (2*N));
            end
        end
    end
end