close all; clear;clc;

load dataDWT.mat;
%load dataAR.mat;
%load dataCSP.mat;
%load dataPSD.mat;
Y(Y==2) = -1;

y=Y;

num_iterations=30;
base_algorithm= 'regression_tree';
min_parent = floor(size(X, 1) / 5) * 4 - 1;
loss='logistic';
learning_rate=1;

no_objects = length(y);
fifth = floor(no_objects / 5);
error_train = zeros([num_iterations 1]);
error_test = zeros([num_iterations 1]);

for fold = 1 : 5
    mask_test = zeros([1 no_objects]);
    mask_test((fold - 1) * fifth + 1 : fifth * fold) = 1;
    mask_test = logical(mask_test);
    train_set = X(~mask_test,:);
    train_ans = y(~mask_test);
    test_set = X(mask_test,:);
    test_ans = y(mask_test);
    model = gradient_boosting_train(train_set, train_ans, num_iterations, ...
        base_algorithm, loss, 'learning_rate', learning_rate, ...
        'min_parent', min_parent);
    [~, error_train_loop] = gradient_boosting_predict(model, train_set, train_ans);
    [~, error_test_loop] = gradient_boosting_predict(model, test_set, test_ans);
    error_train = error_train + error_train_loop;
    error_test = error_test + error_test_loop;
end
error_train = error_train / 5;
error_test = error_test / 5;

hold on;
xlabel('Number of iterations');
ylabel('Error');
title('Gradient boosting - PSD');
plot(1 : num_iterations, error_train, 'linewidth', 2);