%% TEST 3: Custom Hyperparameter Optimization (Grid Search) na zbiorze Ionosphere z użyciem custom training loop

% Ustawienia początkowe
rng('default');

%% 1. Wczytanie i przygotowanie danych
load ionosphere
if ~iscategorical(Y)
    Y = categorical(Y);
end

n = size(X,1);
idx = randperm(n);

% Podział danych: 60% trening, 20% walidacja, 20% test
trainRatio = 0.6;
valRatio   = 0.2;
testRatio  = 0.2;

nTrain = round(trainRatio * n);
nVal   = round(valRatio * n);

trainIdx = idx(1:nTrain);
valIdx   = idx(nTrain+1:nTrain+nVal);
testIdx  = idx(nTrain+nVal+1:end);

XTrain = X(trainIdx, :);
YTrain = Y(trainIdx, :);

XVal   = X(valIdx, :);
YVal   = Y(valIdx, :);

XTest  = X(testIdx, :);
YTest  = Y(testIdx, :);

% Pobranie nazw klas z danych treningowych
classNames = categories(YTrain);

%% 2. Definicja architektury sieci MLP
inputSize = size(X,2);            % 34 cechy
numClasses = numel(classNames);     % liczba klas

layers = [
    featureInputLayer(inputSize, 'Name', 'input')
    fullyConnectedLayer(10, 'Name', 'fc1')
    reluLayer('Name', 'relu1')
    fullyConnectedLayer(numClasses, 'Name', 'fc2')
    softmaxLayer('Name', 'softmax')
    % classificationLayer - nie potrzebujemy, bo obliczamy stratę ręcznie
];

%% 3. Definicja siatki hiperparametrów (grid search)
candidateLearnRates = [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1];
candidateBatchSizes = [8, 16, 32];

bestValError = Inf;
bestLearnRate = NaN;
bestBatchSize = NaN;

numEpochs = 50;  % liczba epok do treningu dla każdej kombinacji

fprintf('Rozpoczynamy optymalizację hiperparametrów...\n');
for lr = candidateLearnRates
    for bs = candidateBatchSizes
        hyperparams.LearnRate = lr;
        hyperparams.BatchSize = bs;
        % Funkcja trainAndValidate wykorzystuje custom training loop i zwraca 1 - accuracy na zbiorze walidacyjnym
        currentValError = trainAndValidate(hyperparams, XTrain, YTrain, XVal, YVal, layers, classNames, numEpochs);
        fprintf('LearnRate: %.5f, BatchSize: %d, Validation Error: %.4f\n', lr, bs, currentValError);
        if currentValError < bestValError
            bestValError = currentValError;
            bestLearnRate = lr;
            bestBatchSize = bs;
        end
    end
end

fprintf('Najlepszy współczynnik uczenia: %.5f\n', bestLearnRate);
fprintf('Najlepszy rozmiar batcha: %d\n', bestBatchSize);

%% 4. Trenowanie finalnej sieci z wykorzystaniem optymalnych hiperparametrów (custom training loop)
finalNet = finalTrain(XTrain, YTrain, layers, bestLearnRate, bestBatchSize, numEpochs, classNames);

%% 5. Ocena finalnej sieci na zbiorze testowym
YPredTest = predictLabels(finalNet, XTest, classNames);
accuracyTest = mean(YPredTest == YTest);
fprintf('Dokładność na zbiorze testowym: %.2f%%\n', accuracyTest * 100);

% Rysowanie macierzy pomyłek
figure;
plotconfusion(YTest, YPredTest);
title('Macierz pomyłek: Finalna sieć MLP (IONOSPHERE)');

%% Wykres 3D punktów (cechy 3, 4 i 5) – nakładanie danych uczących i testowych

% Obliczamy predykcje dla zbioru treningowego i testowego
YPredTrain = classify(net, XTrain);
YPredTest  = classify(net, XTest);

% Określamy, które próbki zostały poprawnie sklasyfikowane
idxTrainCorrect = (YPredTrain == YTrain);
idxTrainIncorrect = ~idxTrainCorrect;

idxTestCorrect = (YPredTest == YTest);
idxTestIncorrect = ~idxTestCorrect;

% Tworzymy wykres 3D
figure;
hold on;
title('Wykres 3D: Dane uczące i testowe (cechy 3,4,5)');
xlabel('Cecha 3');
ylabel('Cecha 4');
zlabel('Cecha 5');

% --- Dane uczące ---
% Rysujemy dane uczące markerami "o" (okręgi):
% Poprawnie sklasyfikowane dane uczące – zielone, wypełnione
scatter3(XTrain(idxTrainCorrect,3), XTrain(idxTrainCorrect,4), XTrain(idxTrainCorrect,5), ...
    50, 'g', 'o', 'filled', 'DisplayName', 'Train Correct');
% Błędnie sklasyfikowane dane uczące – czerwone, okręgi bez wypełnienia
scatter3(XTrain(idxTrainIncorrect,3), XTrain(idxTrainIncorrect,4), XTrain(idxTrainIncorrect,5), ...
    50, 'r', 'o', 'DisplayName', 'Train Incorrect');

% --- Dane testowe ---
% Rysujemy dane testowe markerami "s" (kwadraty):
% Poprawne predykcje – zielone, wypełnione
scatter3(XTest(idxTestCorrect,3), XTest(idxTestCorrect,4), XTest(idxTestCorrect,5), ...
    80, 'b', 's', 'filled', 'DisplayName', 'Test Correct');
% Błędne predykcje – czerwone, kwadraty bez wypełnienia
scatter3(XTest(idxTestIncorrect,3), XTest(idxTestIncorrect,4), XTest(idxTestIncorrect,5), ...
    80, 'k', 's', 'DisplayName', 'Test Incorrect');

grid on;
legend('Location','best');
hold off;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Funkcje pomocnicze

%% Funkcja trainAndValidate - custom training loop (używana do optymalizacji hiperparametrów)
function valError = trainAndValidate(hyperparams, XTrain, YTrain, XVal, YVal, layers, classNames, numEpochs)
    % Konwersja architektury do dlnetwork
    dlnet = dlnetwork(layers);
    
    miniBatchSize = hyperparams.BatchSize;
    learnRate = hyperparams.LearnRate;
    trailingAvg = [];
    trailingAvgSq = [];
    iteration = 0;
    numObservations = size(XTrain,1);
    
    % Pętla treningowa
    for epoch = 1:numEpochs
        idx = randperm(numObservations);
        XTrain = XTrain(idx,:);
        YTrain = YTrain(idx,:);
        
        for i = 1:miniBatchSize:numObservations
            iteration = iteration + 1;
            idxBatch = i:min(i+miniBatchSize-1, numObservations);
            % Przygotowanie danych: transpozycja - format [cechy, batch]
            XBatch = XTrain(idxBatch,:)';
            dlX = dlarray(single(XBatch), 'CB');
            % Konwersja etykiet do postaci one-hot (format: [numClasses, batch])
            T = onehotencode(YTrain(idxBatch,:), 2)';
            % Obliczenie gradientów i straty
            [loss, gradients] = dlfeval(@modelGradients, dlnet, dlX, T);
            % Aktualizacja wag za pomocą Adam
            [dlnet, trailingAvg, trailingAvgSq] = adamupdate(dlnet, gradients, trailingAvg, trailingAvgSq, iteration, learnRate);
        end
    end
    
    % Ewaluacja na zbiorze walidacyjnym
    YPredVal = predictLabels(dlnet, XVal, classNames);
    accuracyVal = mean(YPredVal == YVal);
    valError = 1 - accuracyVal;
end

%% Funkcja finalTrain - finalna iteracja treningu z wykorzystaniem custom training loop
function dlnet = finalTrain(XTrain, YTrain, layers, learnRate, miniBatchSize, numEpochs, classNames)
    dlnet = dlnetwork(layers);
    trailingAvg = [];
    trailingAvgSq = [];
    iteration = 0;
    numObservations = size(XTrain,1);
    
    for epoch = 1:numEpochs
        idx = randperm(numObservations);
        XTrain = XTrain(idx,:);
        YTrain = YTrain(idx,:);
        
        for i = 1:miniBatchSize:numObservations
            iteration = iteration + 1;
            idxBatch = i:min(i+miniBatchSize-1, numObservations);
            XBatch = XTrain(idxBatch,:)';
            dlX = dlarray(single(XBatch), 'CB');
            T = onehotencode(YTrain(idxBatch,:), 2)';
            [loss, gradients] = dlfeval(@modelGradients, dlnet, dlX, T);
            [dlnet, trailingAvg, trailingAvgSq] = adamupdate(dlnet, gradients, trailingAvg, trailingAvgSq, iteration, learnRate);
        end
        
        % Opcjonalnie: wyświetlanie postępu
        fprintf('Epoch %d, Loss: %.4f\n', epoch, double(loss));
    end
end

%% Funkcja predictLabels - predykcja etykiet przy użyciu dlnet
function YPred = predictLabels(dlnet, X, classNames)
    dlX = dlarray(single(X'),'CB');
    dlYPred = forward(dlnet, dlX);
    [~, idx] = max(extractdata(dlYPred), [], 1);
    YPred = categorical(classNames(idx));
    YPred = YPred(:);
end

%% Funkcja modelGradients - oblicza stratę i gradienty
function [loss, gradients] = modelGradients(dlnet, dlX, T)
    dlYPred = forward(dlnet, dlX);
    loss = crossentropy(dlYPred, T, 'TargetCategories', 'independent');
    gradients = dlgradient(loss, dlnet.Learnables);
end