clear all; close all; clc; rng(1);

%
%% Part 1 - Bootstrapping
% In a medical study, doctors aimed to examine whether aspirin can reduce the probability of a heart attack. To this aim, they recruited a large # of subjects and divided them into an aspirin group and placebo group. After receiving the aspirin or the placebo pill for a period of time, doctors observed that people in the aspirin group ended up only having 104 heart attacks, and people in the placebo group ended up having 189.
% Given the data, we can compute the percentage of heart attacks for both groups, and compute a ratio. The fact that the ratio is smaller than 1 suggests that aspirin therapy is effective in preventing heart attacks. But how sure are we of this estimate? Does the 95% confidence interval include 1?

%aspirin group
aspirin_heart   = 104;
aspirin_total   = 11037;
aspirin_data    = [ones(1,aspirin_heart), zeros(1, aspirin_total-aspirin_heart)];
%placebo group
placebo_heart   = 189;
placebo_total   = 11034;
placebo_data    = [ones(1,placebo_heart), zeros(1, placebo_total-placebo_heart)];
% Calculate statistic for original sample
ratio_empirical = (aspirin_heart/aspirin_total)/(placebo_heart/placebo_total);
disp(['The ratio computed from empirical data is ', num2str(ratio_empirical)])

%Let's say we bootstrap (sample with replacement) 10,000 times
n_boot = 1e4;
%ratio_boot is used to store the ratio for each bootstrapped dataset
ratio_boot = zeros(n_boot, 1);

%------------------------YOUR CODE STARTS HERE-------------------------



% ---------------------------------------------------------------------

% Find 95% confidence interval
ratio_boot_sorted = sort(ratio_boot); % Sort ratio_boots from lowest to highest
lower_bound_ind = ceil(0.025 * n_boot); % Find index of 2.5% value
upper_bound_ind = floor(0.975 * n_boot); % Find index of 97.5% value

% Use lower/higher index to find value corresponding to 2.5%/97.5% position
lower_bound = ratio_boot_sorted(lower_bound_ind);
upper_bound = ratio_boot_sorted(upper_bound_ind);
%----------------------------------------------------------------------

%
% Let's visualize the bootstrapped ratios and the 95% confidence interval.

%This section does not require changes.
cMAP1  = [50,205,50; 255,165,0]./255;
figure
h(1) = histogram(ratio_boot, 'FaceColor',cMAP1(1,:), 'FaceAlpha', 0.3, ...
    'EdgeColor', cMAP1(1,:)); hold on
h(2) = plot([1, 1]*ratio_empirical, [0, 1]*700,'-', 'lineWidth',3);
h(3) = plot([1, 1]*lower_bound, [0, 1]*700,'--', 'Color', 'k', 'lineWidth',3);
h(4) = plot([1, 1]*upper_bound, [0, 1]*700,'-.', 'Color', 'k', 'lineWidth',3);
hold off; xlabel('Ratio'); ylabel('Counts'); box off
yticks(0:300:600); xticks(round([lower_bound, ratio_empirical, upper_bound],2));
legend(h, {'Bootstrapped ratios', 'Empirical ratio', '95% CI (lower bound)', ...
    '95 % CI (upper bound)'});
set(gca,'FontSize', 15);
set(gcf, 'Units', 'Normalized', 'OuterPosition', [0, 0.04, 0.4, 0.4]);
set(gcf,'PaperUnits','centimeters','PaperSize',[25 15]);

% The fact that the 95% confidence interval does not include 1 implies that we can be very confident that the aspirin therapy is effective in preventing heart attacks.
%
%% Part 2 - Cross validation (Leave-one-out)
% Cross-validation is a resampling method that uses different portions of the data to test and train a model on different iterations. The goal of cross-validation is to test the model's ability to predict new data that was not used in estimating it, in order to flag problems like overfitting and to give an insight on how the model will generalize to a new (or unknown) dataset.

%let's first load the file
load regress1.mat %x: x-values; y: y-values
%visualize the data
figure
scatter(x, y, 100, 'filled', 'MarkerFaceColor', cMAP1(1,:),...
    'MarkerFaceAlpha',0.5);
xlabel('x'); ylabel('y'); xlim([-2,3]); ylim([-3,4]); grid on
set(gca,'FontSize',15);
set(gcf, 'Units', 'Normalized', 'OuterPosition', [0, 0.04, 0.3, 0.45]);
set(gcf,'PaperUnits','centimeters','PaperSize',[20 15]);

% Now, we will fit the data with polynomial linear models without using cross validation. This will serve as a comparison to see how cross validation alters our interpretation of the 'best' model.

% Number of data points
n_pts    = length(y);
% All regressors that we will use for each fit
XX       = [ones(n_pts,1), x, x.^2,x.^3,x.^4,x.^5,x.^6,x.^7,x.^8,x.^9,x.^10,x.^11];
% Number of regressors (should be 12)
n_models = size(XX,2);
order    = 0:(n_models-1);
%initialize matrices mse_vec (which stores mean squared error) and fit_mat
%(which stores predicted y values)
mse_vec  = zeros(n_models, 1);
fit_mat  = zeros(n_models, n_pts);
betas    = cell(1,n_models);

%--------------------------YOUR CODE STARTS HERE---------------------------




%--------------------------------------------------------------------------

% Visualize fits - which one looks best?
figure
for i = 1:n_models
    subplot(3,4,i)
    scatter(x, y, 100, 'filled', 'MarkerFaceColor', cMAP1(1,:),...
        'MarkerFaceAlpha',0.5); hold on;
    plot(x, fit_mat(i, :), 'LineWidth', 2, 'Color',cMAP1(2,:)); hold on
    xlabel('x'); ylabel('y'); xlim([-2,3]); ylim([-3,4]); grid on
    title_str = append('Order: ', int2str(i-1)); title(title_str)
    set(gca,'FontSize',15);
end
set(gcf, 'Units', 'Normalized', 'OuterPosition', [0, 0.04, 0.7, 0.8]);
set(gcf,'PaperUnits','centimeters','PaperSize',[35 25]);

%
% Next let's re-do the model fitting with leave-one-out cross validation and then re-evaluate our models.

%------------------------YOUR CODE STARTS HERE-----------------------------



%--------------------------------------------------------------------------

% Plot the mean MSE
figure
plot(order, mean_mse_test, 'linewidth',2, 'Marker', 's', 'Color',...
    cMAP1(2,:), 'MarkerSize',15); hold on
scatter(min_ind-1, min_mse,150,'rs','filled'); hold on
plot(order, mean_mes_train , '-*','linewidth',2, 'Color',...
    cMAP1(1,:), 'MarkerSize',6); hold off
xlim([0,n_models-1] + [-0.5, 0.5]); xlabel('Order'); grid on;
yticks(0:1:5); ylabel('Mean squared error (MSE)'); box off;
legend({'MSE (test set)','MSE of the ''best'' model', 'MSE (training set)'});
set(gca,'FontSize',15);
set(gcf, 'Units', 'Normalized', 'OuterPosition', [0, 0.04, 0.3, 0.7]);
set(gcf,'PaperUnits','centimeters','PaperSize',[25 30]);
%saveas(gcf, 'Fig3', 'pdf');

% Which model is best?
disp(['The best model is order = ', num2str(order(min_ind))]);

%
%% Part 3 - Regularization
% In this section, we will use Ridge (L2) regularization to reduce variance of an overparameterized model.
% First we need data generated by a groundtruth function. Let's say this function is a 7th degree polynomial. So the points we will generate lie on a 7-degree polynomial function with some added noise.
% Now, let's assume our model has 30 degrees of freedom (akin to a 29 degree polynomial). We can find the best fit coeffs by using np.polyfit (built-in function)
% The 29 degree polynomial gives us a very bad fit to the held-out test data. Let's regularize with ridge and lasso and observe how test error changes

% Generate train data
N = 300;
x = linspace(-1, 1, N)';
y = (10 * x.^7 - 8 * x.^5 + 3 * x.^3 - 2 * x + 0.5) + randn(N,1);

% Divide into train and test sets

% Overfit with a high-degree polynomial

% Define regularization parameters

% Train models with Ridge regularization

% Predict for both models

% Plot results
figure;
plot(x, y, 'o', 'DisplayName', 'Train data');
hold on;
plot(x, y_pred_poly, 'DisplayName', sprintf('Polynomial (degree=%d)', degree));
plot(x, y_pred_ridge, 'DisplayName', sprintf('Ridge (alpha=%.2f)', ridge_alpha));
xlabel('X');
ylabel('Y');
legend;
title('Comparison of models');
grid on;

%

% Predict for all models on test data

% Plot results
figure;
plot(x_test, y_test, 'o', 'DisplayName', 'Test data');
hold on;
plot(x_test, y_pred_poly_test, 'DisplayName', sprintf('Polynomial (degree=%d)', degree));
plot(x_test, y_pred_ridge_test, 'DisplayName', sprintf('Ridge (alpha=%.2f)', ridge_alpha));
xlabel('X');
ylabel('Y');
legend;
ylim([-100, 300]);
title('Comparison of models');
grid on;

%

% Calculate mean squared error (MSE)
mse_poly = mean((y_pred_poly_test - y_test).^2);
mse_ridge = mean((y_pred_ridge_test - y_test).^2);

% Plot MSE values
figure;
bar([mse_poly, mse_ridge]);
set(gca, 'xticklabel', {'Polynomial', 'Ridge'});
xlabel('Model');
ylabel('Mean Squared Error (MSE)');
set(gca, 'YScale', 'log');
title('Comparison of MSE on held-out data');

%

% Extract coefficients

% Plot coefficients
figure;

% Polynomial model
subplot(3,1,1);
stem(0:degree, coeffs_poly);
xlabel('Polynomial degree');
ylabel('Coefficient value');
title('Polynomial Coefficients');
grid on;

% Ridge model
subplot(3,1,2);
stem(0:degree, coeffs_ridge);
xlabel('Feature index');
ylabel('Coefficient value');
title(['Ridge Coefficients (alpha=', num2str(ridge_alpha), ')']);
grid on;

%
%% Part 4 - Permutation Testing
% In Master, Subramanian, & Zhao, 2005, Science, the authors ran the above task and computational models of behavior. After confirming accurate model recoverability with a generate & recover procedure, they found a significant correlation between Math Tools grade & attentional slope. They concluded that students who do well in Math Tools are better at allocating their attention to relevant information.
% Let's try and confirm their finding using a permutation test. The r-value they published in the paper is 0.23.

% First, let's load the original data from their paper.
load Master_et_al_2005.mat

nsubjs = 30;
r_stat = 0.23;

% At the core of a permutation test is the creation of a null distribution
% of results. I.e. How many times do we achieve the r-value Master et al.
% published, purely due to chance?

% To do this we permute (randomize) our variables of interest such that we 
% no there is not any relationship between them. 
% Then we compare the published r-value to our null distribution to 
% compute a p-statistic (a percentage describing how many times our null
% distribution achieves the value originally obtained).

n_perms = 1000;
r_perms = NaN(n_perms,1);

%---------------------------- YOUR CODE HERE -------------------------%

p_value = [];

%---------------------------- YOUR CODE HERE -------------------------%

% Find confidence intervals
ratio_perm_sorted = sort(r_perms); 
lower_bound_ind = ceil(0.025 * n_perms);
upper_bound_ind = floor(0.975 * n_perms); 

% Use lower/higher index to find value corresponding to 2.5%/97.5% position
lower_bound = ratio_perm_sorted(lower_bound_ind);
upper_bound = ratio_perm_sorted(upper_bound_ind);

% Let's see whether we can conclude with confidence that there is a
% relationship between attentional focus & math tools grades.

figure
histogram(r_perms)
hold on
plot([r_stat,r_stat],ylim,'--r','LineWidth',2)
plot([lower_bound,lower_bound],ylim,'--k')
plot([upper_bound,upper_bound],ylim,'--k')
legend({'Null dist','Original r-statistic','Lower 95% conf. bound','Upper 95% conf. bound'})
title(['p = ' num2str(p_value)])