% Load the MNIST dataset, rescaled
% Nutno stahnout mnist.mat -> https://marian.fsik.cvut.cz/~musil/ZAPG/Matlab/mnist.mat
load ('mnist.mat')

% training data 
X_train = training.images(:,:,:);
Y_train = training.labels(:);
% test data
X_test = test.images(:,:,:);
Y_test = test.labels(:);

% Reshape data for processing
N_pix   = size(X_train,1)*size(X_train,2);
N_train = size(X_train,3);
N_test  = size(X_test,3);
X_train = reshape(X_train, N_pix, N_train);
X_test  = reshape(X_test,  N_pix, N_test);

% Y_train(10:15)

% Convert the labels to one-hot encoding
Y_train = [Y_train, zeros(size(Y_train,1),9)]';
for i = 1:size(Y_train,2)
    Y = zeros(10,1);
    Y(Y_train(1,i)+1) = 1;
    Y_train(:,i) = Y;
end
Y_test = [Y_test, zeros(size(Y_test,1),9)]';
for i = 1:size(Y_test,2)
    Y = zeros(10,1);
    Y(Y_test(1,i)+1) = 1;
    Y_test(:,i) = Y;
end

disp('Loading and formating data, done.')


% User-defined parameters
input_layer_size = size(X_train, 1); % Input layer size (784 for MNIST)
hidden_layers = [80,30];             % User-defined hidden layer sizes
output_layer_size = 10;              % Output layer size (10 for MNIST digits)
% Training parameters
learning_rate = 4;
epochs = 30;
batch_size = 20;

% Initialize weights and biases for hidden layers and output layer
L = length(hidden_layers) + 1;
W = cell(L, 1);
b = cell(L, 1);
a = cell(L, 1);
z = cell(L, 1);
d = cell(L, 1);
db = cell(L, 1);
dW = cell(L, 1);

layers = [input_layer_size, hidden_layers, output_layer_size];

for i = 1:L
    W{i}  = randn(layers(i+1), layers(i)) * sqrt(2 / layers(i)); % He initialization
    b{i}  = randn(layers(i+1), 1);
    a{i}  = zeros(layers(i+1), 1);
    z{i}  = zeros(layers(i+1), 1);
    d{i}  = zeros(layers(i+1), 1);
end

% Training loop
for epoch = 1:epochs

    % Shuffle training data
    shuffled_indices = randperm(size(X_train, 2));
    X_train = X_train(:,shuffled_indices);
    Y_train = Y_train(:,shuffled_indices);

    % Mini-batch gradient descent
    for batch = 1:batch_size:size(X_train, 1)
        
        batch_end = min(batch + batch_size - 1, size(X_train, 1));
        
        X_batch = X_train(:,batch:batch_end);
        Y_batch = Y_train(:,batch:batch_end);

        M = size(X_batch,2);
        
        for i = 1:L
            dW{i}  = zeros(layers(i+1), layers(i));
            db{i} =  zeros(layers(i+1), 1);
        end

        for m=1:M

            % forward propagation
            z{1} = W{1} * X_batch(:,m) + b{1}; % Input layer
            a{1} = sig(z{1});
            for i = 2:L
                z{i} = W{i} * a{i-1} + b{i};
                a{i} = sig(z{i});
            end

            % backpropagation
            d{L} = (a{L}-Y_batch(:,m)) .* dsig(z{L});
            for i=L-1:-1:1
                d{i} = ( W{i+1}' * d{i+1} ) .* dsig(z{i});
            end

            % add errors
            db{1} = db{1} + d{1};
            dW{1} = dW{1} + (d{1} * X_batch(:,m)');
            for i = 2:L
                dW{i} = dW{i} + (d{i} * a{i-1}');
                db{i} = db{i} + d{i};
            end

        end
        % update weights and biases
        for i = 1:L
            W{i} = W{i} - learning_rate * (1/M) * dW{i};
            b{i} = b{i} - learning_rate * (1/M) * db{i};
        end

    end

    %forward propagation on test data
    K = 0;
    for j=1:size(Y_test,2)
        z{1} = W{1} * X_test(:,j) + b{1}; % Input layer
        a{1} = sig(z{1});
        for i = 2:L
            z{i} = W{i} * a{i-1} + b{i};
            a{i} = sig(z{i});
        end
        [~, idx_X] = max(a{L});
        [~, idx_Y] = max(Y_test(:,j));
        if idx_Y == idx_X
            K = K + 1;
        end
    end
    disp(['Epoch: ', num2str(epoch), ', Accuracy: ', num2str(K),' / 10 000']);
end

%%%% FUNCTIONS %%%%
function sig = sig(z)
    sig = 1.0 ./ (1.0 + exp(-z));
end

function dsig = dsig(z)
    dsig = sig(z) .* (1 - sig(z));
end

function Z = ReLU(Z)
    Z = max(Z, 0);
end

function dZ = dReLU(Z)
    dZ = Z > 0;
end


function fr = elu(x)
    f = zeros(length(x),1);
    for i = 1:length(x)
    if x(i)>=0
        f(i) = x(i);
    else
        f(i) = 0.2*(exp(x(i))-1);
    end
    end
    fr = f;
end

function fr = elup(x)
    f = zeros(length(x),1);
    for i = 1:length(x)
    if x(i)>=0
        f(i) = 1;
    else
        f(i) = 0.2*exp(x(i));
    end
    end
    fr = f;
end

