% ============================================================
% Image Processing + Machine Learning 
% ============================================================

clc; clear; close all;

%% ============================================================
% STEP 1 — Load Dataset
% ============================================================

dataDir = fullfile(pwd,'data');

imds = imageDatastore(dataDir, ...
    'IncludeSubfolders', true, ...
    'LabelSource','foldernames');

disp('Class distribution:');
disp(countEachLabel(imds));

%% ============================================================
% STEP 2 — Load ONE image for demonstration
% ============================================================

testImagePath = fullfile(pwd,'data','example.jpg');

if ~isfile(testImagePath)
    error('Place an image named example.jpg inside the data folder.');
end

I = imread(testImagePath);

figure('Name','Step 2: Original Image');
imshow(I);
title('Original Image');

%% ============================================================
% STEP 3 — Convert to Grayscale
% ============================================================

if ndims(I) == 3
    Igray = rgb2gray(I);
else
    Igray = I;
end

figure('Name','Step 3: Grayscale');
subplot(1,2,1); imshow(Igray); title('Grayscale Image');
subplot(1,2,2); imhist(Igray); title('Histogram (Grayscale)');

%% ============================================================
% STEP 4 — Resize Image
% ============================================================

Igray = imresize(Igray, [128 128]);

figure('Name','Step 4: Resized Image');
subplot(1,2,1); imshow(Igray); title('Resized (128x128)');
subplot(1,2,2); imhist(Igray); title('Histogram (Resized)');

%% ============================================================
% STEP 5 — Noise Reduction (Gaussian Filtering)
% ============================================================

I_denoised = imgaussfilt(Igray, 1);

figure('Name','Step 5: Denoising');
subplot(1,2,1); imshow(I_denoised); title('Denoised Image');
subplot(1,2,2); imhist(I_denoised); title('Histogram (Denoised)');

%% ============================================================
% STEP 6 — Contrast Enhancement
% ============================================================

I_histeq = histeq(Igray);
I_adapthisteq = adapthisteq(Igray);

figure('Name','Step 6: Contrast Enhancement');
subplot(2,2,1); imshow(Igray); title('Original Gray');
subplot(2,2,2); imhist(Igray); title('Histogram');
subplot(2,2,3); imshow(I_histeq); title('Histogram Equalization');
subplot(2,2,4); imshow(I_adapthisteq); title('Adaptive HistEq');

%% ============================================================
% STEP 7 — Edge Detection
% ============================================================

I_edges = edge(Igray, 'Canny');

figure('Name','Step 7: Edge Detection');
subplot(1,2,1); imshow(I_edges); title('Edges (Canny)');
subplot(1,2,2); imhist(I_edges); title('Histogram (Edges)');

%% ============================================================
% STEP 8 — Feature Extraction (for ML)
% ============================================================
% Features:
%   - Mean intensity
%   - Standard deviation
%   - Entropy

feat_mean = mean(I_denoised(:));
feat_std  = std(double(I_denoised(:)));
feat_ent  = entropy(I_denoised);

figure('Name','Step 8: Feature Values');
bar([feat_mean feat_std feat_ent]);
set(gca,'XTickLabel',{'Mean','Std','Entropy'});
title('Extracted Numerical Features');

%% ============================================================
% STEP 9 — Build Feature Dataset from ALL Images (ML Training)
% ============================================================

features = [];
labels   = imds.Labels;

for i = 1:numel(imds.Files)

    Itemp = readimage(imds, i);

    if ndims(Itemp) == 3
        Itemp = rgb2gray(Itemp);
    end

    Itemp = imresize(Itemp, [128 128]);
    Itemp = imgaussfilt(Itemp, 1);

    f = [
        mean(Itemp(:))
        std(double(Itemp(:)))
        entropy(Itemp)
    ];

    features = [features; f'];
end

%% ============================================================
% STEP 10 — Train Machine Learning Model (kNN)
% ============================================================

k = 3;
mdl = fitcknn(features, labels, 'NumNeighbors', k);

disp('kNN model trained.');

%% ============================================================
% STEP 11 — Predict Class of the Test Image
% ============================================================

testFeat = [feat_mean feat_std feat_ent];
predictedLabel = predict(mdl, testFeat);

figure('Name','Step 11: Prediction Result');
text(0.1,0.5,"Predicted Class: " + string(predictedLabel),'FontSize',14);
axis off;

%% ============================================================
% STEP 12 — Feature Space Visualization WITH Test Image
% ============================================================
% This plot shows:
% - Training images as colored points (by class)
% - The example image as a black star
% ============================================================

figure('Name','Feature Space with Test Image','NumberTitle','off');
hold on;

% Plot training images
gscatter(features(:,1), features(:,2), labels);

% Plot test image (mean vs std)
plot(feat_mean, feat_std, 'kp', ...
     'MarkerSize', 14, ...
     'MarkerFaceColor', 'k');

% Annotate test image
text(feat_mean, feat_std, ...
     '  Test Image', ...
     'FontWeight','bold');

xlabel('Mean Intensity');
ylabel('Standard Deviation');
title('Feature Space (Training Images + Test Image)');
grid on;
legend('Location','best');
hold off;