%% clear all variables and start fresh
clear variables

%% set seed
seed = 5;
rng(seed)

%% problem parameters

n = 1000;
gamma = 2;
p = round(n * gamma);

n_test = 1e5;

%% generate train and test features

feat_cov_type = 'iso';
feat_cov_param = 0.25;
feat_cov = generate_feature_covariance(p, feat_cov_type, feat_cov_param);

feat_type = 'gau';
feat_param = 5;

X = generate_feature_matrix(n, p, feat_cov, feat_type, feat_param);
X_test = generate_feature_matrix(n_test, p, feat_cov, feat_type, feat_param);

%% generate train and test responses

% % signal generation
% sig_type = 'iso';
% sig_energy = 1;
% beta0 = sqrt(sig_energy) * generate_signal_vector(p, sig_type);

% [W, R] = eig(feat_cov); % eigenvalue decomposition, R: eigenvalues
% % reorder the eigenvalues and eigenvectors of Sigma
% [~, ind] = sort(diag(R), 'descend'); % get the sort indices
% R = R(ind, ind); % r_1 largest
% W = W(:, ind); % reorder the eigenvectors accordingly
% dR = diag(R); % extract eigenvalues
% sig_type = 'top';
% sig_energy = 50;
% beta0 = W(:, 1) / sqrt(dR(1)) * sqrt(sig_energy);

% signal generation
sig_type = 'iso';
sig_energy = 1;
beta0 = sqrt(sig_energy) * generate_signal_vector(p, sig_type);

% noise generation
noi_type = 'gau';
noi_param = 5;
noi_energy = 1;
eps = sqrt(noi_energy) * generate_noise_vector(n, noi_type, noi_param);
eps_test = sqrt(noi_energy) * generate_noise_vector(n_test, noi_type, noi_param);

% response generation
mod_type = 'lin';
y = generate_response_vector(X, beta0, eps, mod_type, feat_cov);
y_test = generate_response_vector(X_test, beta0, eps_test, mod_type, feat_cov);

%% precomputation for GCV

S_gram = (X * X') / n; % sample gram matrix of dimension n \times n

[U, D] = eig(S_gram); % eigenvalue decomposition of S_gram
s_vals = diag(D); % eigenvalues in a vector

%% loocv

% X: data matrix of dimension n times p
% y: response vector of dimension n
% X_test: test data matrix of dimension n_test times p
% y_test: test response vector of dimension n_test

n = size(X, 1); % number of data points
p = size(X, 2); % number of features
n_test = size(X_test, 1); % number of test data points

% Define gradient descent hyperparameters
alpha = 0.01; % learning rate
num_iters = 1000; % number of iterations

% Initialize leave-one-out cross-validation score array
loo_scores = zeros(num_iters, n);

% Initialize leave-one-out residual array
loo_resids = zeros(num_iters, n);

% Initialize test error array
test_errors = zeros(num_iters, 1);

% Initialize gcv residual array
gcv_resids = zeros(num_iters, n);
gcv_scores = zeros(num_iters, n);


% % Perform leave-one-out cross-validation
% for k = 1:n
%     disp(k)
%     % Define training and validation sets
%     X_train = X([1:k-1 k+1:n], :);
%     y_train = y([1:k-1 k+1:n]);
%     X_val = X(k, :);
%     y_val = y(k);
%     
%     % Initialize theta to zero vector
%     theta = zeros(p, 1);
%     
%     % Perform gradient descent on training set
%     for i = 1:num_iters
%         theta = theta - alpha * X_train' * (X_train * theta - y_train) / n;
%         
%         % Compute the loo residuals
%         loo_resids(i, k) = y_val - X_val * theta;
%         
%         % Compute validation score on validation set
%         loo_scores(i, k) = loo_resids(i, k)^2;
%     end
%     
% end


% Train full model on all data
theta_full = zeros(p, 1);
for i = 1:num_iters
    theta_full = theta_full - alpha * X' * (X * theta_full - y) / n;
    
    disp(i)
    % Compute GCV of full model for all iterations
    gcv_denom = mean(1 - (1 - (1 - alpha * s_vals).^i));
    gcv_resids(i, :) = ((y - X * theta_full) ./ gcv_denom)';
    
    gcv_scores(i, :) = gcv_resids(i, :).^2;
    
    % Compute test error of full model for all iterations
    test_errors(i) = (X_test * theta_full - y_test)' * (X_test * theta_full - y_test) / n_test;
end



%%
% X, y, X_test, y_test, alpha, num_iters, loo_scores, and test_errors are defined as in the previous code block

% Initialize arrays to store the means and standard deviations of the leave-one-out scores and test errors
% loo_means = mean(loo_scores, 2);
% loo_stds = std(loo_scores, [], 2);
test_means = mean(test_errors, 2);
test_stds = std(test_errors, [], 2);

gcv_means = mean(gcv_scores, 2);

%%

font_size_legend = 14;
font_size_axis = 16;

figure;
semilogx(1:num_iters, test_means, ...
    'LineWidth', 2, ...
    'DisplayName', 'Risk');
hold on
grid on
semilogx(1:num_iters, gcv_means, ...
    'LineWidth', 2, ...
    'DisplayName', 'GCV');

legend box off
legend('Location', 'Best', 'FontSize', font_size_legend)
xlabel('Number of iterations ($k$)', 'Interpreter', 'latex', 'FontSize', font_size_axis)
ylabel('Value', 'Interpreter', 'latex', 'FontSize', font_size_axis)
% ylim([-Inf, 10])

%% save plot
save_str = strcat(...
    'inconsistency', ...
    '10gamma', num2str(10 * gamma), '_', ...
    'nt', num2str(n_test), '_', ...
    'featcov', feat_cov_type, '_', ...
    'feat', feat_type, '_', ...
    '10featparam', num2str(10 * feat_param), '_', ...
    'sig', sig_type, '_', ...
    'noi', noi_type, '_', ...
    '10noiparam', num2str(10 * noi_param), '_', ...
    'mod', mod_type, '_', ...
    'seed', num2str(seed), '_', ...
    'n', num2str(n),);

save_fig_str = strcat('../figures/inconsistency/', save_str);
saveas(gcf, save_fig_str)
saveas(gcf, save_fig_str, 'epsc')
saveas(gcf, save_fig_str, 'png')

%%

% %%
% % Plot the means and standard deviations of the leave-one-out scores and test errors as a function of the number of iterations
% figure;
% errorbar(1:num_iters, loo_means, loo_stds);
% hold on;
% errorbar(1:num_iters, test_means, test_stds);
% title('Leave-One-Out Scores and Test Errors vs. Number of Iterations');
% xlabel('Number of Iterations');
% ylabel('Score');
% legend('Leave-One-Out Scores', 'Test Errors');
% 
% %% save plot
% save_str = strcat(...
%     'squaredrisk_', ...
%     '10gamma', num2str(10 * gamma), '_', ...
%     'nt', num2str(n_test), '_', ...
%     'featcov', feat_cov_type, '_', ...
%     'feat', feat_type, '_', ...
%     '10featparam', num2str(10 * feat_param), '_', ...
%     'sig', sig_type, '_', ...
%     'noi', noi_type, '_', ...
%     '10noiparam', num2str(10 * noi_param), '_', ...
%     'mod', mod_type, '_', ...
%     'seed', num2str(seed));
% 
% save_fig_str = strcat('../figures/coverage/', save_str);
% saveas(gcf, save_fig_str)
% saveas(gcf, save_fig_str, 'epsc')
% saveas(gcf, save_fig_str, 'png')
% 
% %%
% 
% %% plots
% marker_list = ['*', 'o', '+', 'x'];
% 
% font_size_legend = 14;
% font_size_axis = 16;
% 
% 
% % interval coverage
% figure
% % hold on
% % grid on
% 
% for j = 1:length(coverage_grid(16:end))
%     marker = marker_list(j);
% %     color = get(groot,'defaultAxesColorOrder');
%     legend_str = ['Coverage level = ', num2str(coverage_grid(15+j))];
%     semilogx(10:num_iters, coverageval_loo_grid(10:end, 15+j), ...
%     'LineWidth', 2, ...
%     'DisplayName', legend_str);
%     hold on
%     yline(coverage_grid(15+j), '--', 'HandleVisibility','off')
% end
% % hold on
% % grid on
% % plot(coverage_grid, coverageval_loo_grid, ...
% %     '*-', ...
% %     'LineWidth', 2, ...
% %     'DisplayName', 'LOOCV');
% % plot_grid = linspace(0, 1);
% % plot(plot_grid, plot_grid, '--', ...
% %     'LineWidth', 1.5, ...
% %     'DisplayName', 'Ideal')
% % legend show
% grid on
% ylim([0.6, 1])
% legend box off
% legend('Location', 'SouthEast', 'FontSize', font_size_legend)
% xlabel('Number of iterations ($k$)', 'Interpreter', 'latex', 'FontSize', font_size_axis)
% ylabel('Prediction interval coverage', 'Interpreter', 'latex', 'FontSize', font_size_axis)
% 
% %% save plot
% save_str = strcat(...
%     'coverage_comp_', ...
%     '10gamma', num2str(10 * gamma), '_', ...
%     'nt', num2str(n_test), '_', ...
%     'featcov', feat_cov_type, '_', ...
%     'feat', feat_type, '_', ...
%     '10featparam', num2str(10 * feat_param), '_', ...
%     'sig', sig_type, '_', ...
%     'noi', noi_type, '_', ...
%     '10noiparam', num2str(10 * noi_param), '_', ...
%     'mod', mod_type, '_', ...
%     'seed', num2str(seed));
% 
% save_fig_str = strcat('../figures/coverage/', save_str);
% saveas(gcf, save_fig_str)
% saveas(gcf, save_fig_str, 'epsc')
% saveas(gcf, save_fig_str, 'png')
% % 
% % interval length
% figure
% % hold on
% % grid on
% for j = 1:length(coverage_grid(16:end))
%     marker = marker_list(j);
%     legend_str = ['Coverage level = ', num2str(coverage_grid(15+j))];
%     semilogx(10:num_iters, predintval_length_loo_grid(10:end, 15+j), ...
%     'LineWidth', 2, ...
%     'DisplayName', legend_str);
%     hold on
% 
% % plot(coverage_grid, predintval_length_loo_grid, ...
% %     '*-', ...
% %     'DisplayName', 'loo');
% end
% grid on
% legend show
% legend box off
% legend('Location', 'NorthEast', 'FontSize', font_size_legend)
% xlabel('Number of iterations ($k$)', 'Interpreter', 'latex', ...
%     'FontSize', font_size_axis)
% ylabel('Prediction interval length', 'Interpreter', 'latex', ...
%     'FontSize', font_size_axis)
% % 
% save_str = strcat(...
%     'length_comp_', ...
%     '10gamma', num2str(10 * gamma), '_', ...
%     'nt', num2str(n_test), '_', ...
%     'featcov', feat_cov_type, '_', ...
%     'feat', feat_type, '_', ...
%     '10featparam', num2str(10 * feat_param), '_', ...
%     'sig', sig_type, '_', ...
%     'noi', noi_type, '_', ...
%     '10noiparam', num2str(10 * noi_param), '_', ...
%     'mod', mod_type, '_', ...
%     'seed', num2str(seed));
% 
% save_fig_str = strcat('../figures/coverage/', save_str);
% saveas(gcf, save_fig_str)
% saveas(gcf, save_fig_str, 'epsc')
% saveas(gcf, save_fig_str, 'png')
% 
% %% save data
% 
% save_data_str = strcat('../data/coverage/', save_str);
% save(save_data_str)