%% clear all variables and start fresh
clear variables

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

%% problem parameters

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

n_test = 1e5;

%% generate train and test features

feat_cov_type = 'equi';
feat_cov_param = 0.1;
feat_cov = generate_feature_covariance(p, feat_cov_type, feat_cov_param);

feat_type = 't';
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 = 'bot';
sig_energy = 50;
beta0 = W(:, p) / sqrt(dR(p)) * sqrt(sig_energy);

%%
% noise generation
noi_type = 't';
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);

%% get the minimum eigenvalue of X X' / n

gram_mat = X * X' / n;

[~, S] = eig(gram_mat);

s_min = min(diag(S));

lambda_min = -1 * s_min + 0.01;
% lambda_min = -0.15;
lambda_max = 2;
n_lambda_grid = 100;
lambda_grid = linspace(lambda_min, lambda_max, n_lambda_grid);

% lambda_grid = [-0.15, 0, 1, 10];
% n_lambda_grid = length(lambda_grid);
%% vector to store coverage and lengths

% coverage_grid = 0.05:.05:0.95; % go from 0.5 to 0.95

coverage_level = 0.95;

% % to store empirical coverage values
% coverageval_gcv_grid = zeros(length(coverage_grid), n_lambda_grid);
% coverageval_loo_grid = zeros(length(coverage_grid), n_lambda_grid);
% 
% % to store length of prediction intervals
% predintval_length_gcv_grid = zeros(length(coverage_grid), n_lambda_grid);
% predintval_length_loo_grid = zeros(length(coverage_grid), n_lambda_grid);
% predintval_length_oos_grid = zeros(length(coverage_grid), n_lambda_grid);

% to store functional values for abs, sq, med of abs

% absolute error
func_gcv_abs_grid = zeros(n_lambda_grid, 1);
func_loo_abs_grid = zeros(n_lambda_grid, 1);
func_oos_abs_grid = zeros(n_lambda_grid, 1);

% % median of the absolute values
% func_gcv_abs_med_grid = zeros(n_lambda_grid, 1);
% func_loo_abs_med_grid = zeros(n_lambda_grid, 1);
% func_oos_abs_med_grid = zeros(n_lambda_grid, 1);

% squared error
func_gcv_sq_grid = zeros(n_lambda_grid, 1);
func_loo_sq_grid = zeros(n_lambda_grid, 1);
func_oos_sq_grid = zeros(n_lambda_grid, 1);

% huber error

func_gcv_huber_grid = zeros(n_lambda_grid, 1);
func_loo_huber_grid = zeros(n_lambda_grid, 1);
func_oos_huber_grid = zeros(n_lambda_grid, 1);

% length of pred intervals

predintval_length_gcv_grid = zeros(n_lambda_grid, 1);
predintval_length_loo_grid = zeros(n_lambda_grid, 1);
predintval_length_oos_grid = zeros(n_lambda_grid, 1);

% threshold for huber loss

thresh = 4;

for j = 1:n_lambda_grid
%% fit estimator and get smoother

% lambda = -1 * 0.1;
% lambda = -1 * 0.1;

disp(j)
lambda = lambda_grid(j);

ridge_estim = fit_ridge(X, y, lambda);

ridge_smoother = get_ridge_smoother(X, lambda);

%% get predicted values and compute residuals

hy = predict_linestim(ridge_estim, X);
resid_orig = y - hy;

hy_test = predict_linestim(ridge_estim, X_test);
resid_oos = y_test - hy_test;

resid_gcv = resid_orig / (1 - trace(ridge_smoother) / n);
resid_loo = resid_orig ./ (1 - diag(ridge_smoother));

if (lambda == 0) && (p > n)
   
    resid_mod = gram_mat \ y;
    
    resid_gcv = resid_mod / (trace(pinv(gram_mat)) / n);
    resid_loo = resid_mod ./ diag(pinv(gram_mat));
    
end

    func_gcv_abs = mean(abs(resid_gcv));
    func_loo_abs = mean(abs(resid_loo));

    func_gcv_abs_grid(j) = func_gcv_abs;
    func_loo_abs_grid(j) = func_loo_abs;
    
    func_gcv_abs_med = median(abs(resid_gcv));
    func_loo_abs_med = median(abs(resid_loo));
    
%     func_gcv_abs_med_grid(j) = func_gcv_abs_med;
%     func_loo_abs_med_grid(j) = func_loo_abs_med;
    
    func_gcv_sq = mean(resid_gcv.^2);
    func_loo_sq = mean(resid_loo.^2);

    func_gcv_sq_grid(j) = func_gcv_sq;
    func_loo_sq_grid(j) = func_loo_sq;
    
    func_gcv_huber = mean(compute_huber_loss(resid_gcv, thresh));
    func_loo_huber = mean(compute_huber_loss(resid_loo, thresh));

    func_gcv_huber_grid(j) = func_gcv_huber;
    func_loo_huber_grid(j) = func_loo_huber;
    

%% oos quantities

    func_oos_abs = mean(abs(resid_oos));
    func_oos_abs_grid(j) = func_oos_abs;
    
%     func_oos_abs_med = median(abs(resid_oos));
%     func_oos_abs_med_grid(j) = func_oos_abs_med;
    
    func_oos_sq = mean(resid_oos.^2);
    func_oos_sq_grid(j) = func_oos_sq;

    func_oos_huber = mean(compute_huber_loss(resid_oos, thresh));
    func_oos_huber_grid(j) = func_oos_huber;
    
%% form prediction intervals and get length of prediction intervals


    quanthigh_gcv_grid = quantile(resid_gcv, 1 - (1 - coverage_level)/2);
    quanthigh_loo_grid = quantile(resid_loo, 1 - (1 - coverage_level)/2);
    quanthigh_oos_grid = quantile(resid_oos, 1 - (1 - coverage_level)/2);

    quantlow_gcv_grid = quantile(resid_gcv, (1- coverage_level)/2);
    quantlow_loo_grid = quantile(resid_loo, (1- coverage_level)/2);
    quantlow_oos_grid = quantile(resid_oos, (1- coverage_level)/2);
    
    predintval_length_gcv_grid(j) = ...
        quanthigh_gcv_grid - quantlow_gcv_grid;
    predintval_length_loo_grid(j) = ...
        quanthigh_loo_grid - quantlow_loo_grid;
    predintval_length_oos_grid(j) = ...
        quanthigh_oos_grid - quantlow_oos_grid;

end

%% tempoary plotting

figure
hold on
grid on
plot(lambda_grid, sqrt(func_oos_sq_grid), ...
    'Color', [0 0.4470 0.7410], ...
    'LineWidth', 1.5, ...
    'DisplayName', 'Squared error');
[~, idx] = min(func_oos_sq_grid);
plot(lambda_grid(idx), sqrt(func_oos_sq_grid(idx)), '*', ...
    'LineWidth', 1.5, ...
    'Color', [0 0.4470 0.7410], 'HandleVisibility', 'off');
plot(lambda_grid, sqrt(func_gcv_sq_grid), ...
    'LineStyle', '-.', ...
    'LineWidth', 1.5, ...
    'Color', [0 0.4470 0.7410], ...
    'HandleVisibility', 'off');
%     'DisplayName', 'GCV squred error');
[~, idx] = min(func_gcv_sq_grid);
plot(lambda_grid(idx), sqrt(func_gcv_sq_grid(idx)), 'o', ...
    'LineWidth', 1.5, ...
    'Color', [0 0.4470 0.7410], 'HandleVisibility', 'off');

plot(lambda_grid, sqrt(func_oos_huber_grid), ...
    'Color', [0.8500 0.3250 0.0980], ...
    'LineWidth', 1.5, ...
    'DisplayName', 'Huber error');
[~, idx] = min(func_oos_huber_grid);
plot(lambda_grid(idx), sqrt(func_oos_huber_grid(idx)), '*', ...
    'LineWidth', 1.5, ...
    'Color', [0.8500 0.3250 0.0980], 'HandleVisibility', 'off');
plot(lambda_grid, sqrt(func_gcv_huber_grid), ...
    'LineStyle', '-.', ...
    'LineWidth', 1.5, ...
    'Color', [0.8500 0.3250 0.0980], ...
    'HandleVisibility', 'off');
%     'DisplayName', 'GCV Huber error');
[~, idx] = min(func_gcv_huber_grid);
plot(lambda_grid(idx), sqrt(func_gcv_huber_grid(idx)), 'o', ...
    'LineWidth', 1.5, ...
    'Color', [0.8500 0.3250 0.0980], 'HandleVisibility', 'off');

plot(lambda_grid, func_oos_abs_grid, ... 
    'LineWidth', 1.5, ...
    'Color', [0.9290 0.6940 0.1250], ...
    'DisplayName', 'Absolute error');
[~, idx] = min(func_oos_abs_grid);
plot(lambda_grid(idx), func_oos_abs_grid(idx), '*', ...
    'LineWidth', 1.5, ...
    'Color', [0.9290 0.6940 0.1250], 'HandleVisibility', 'off');
plot(lambda_grid, func_gcv_abs_grid, ...
    'LineWidth', 1.5, ...
    'LineStyle', '-.', ...
    'Color', [0.9290 0.6940 0.1250], ...
    'HandleVisibility', 'off');
%     'DisplayName', 'GCV absolute error');
[~, idx] = min(func_gcv_abs_grid);
plot(lambda_grid(idx), func_gcv_abs_grid(idx), 'o', ...
    'LineWidth', 1.5, ...
    'Color', [0.9290 0.6940 0.1250], 'HandleVisibility', 'off');

% plot(lambda_grid, func_oos_abs_med_grid, ...
%     'DisplayName', 'Median absolute error');
% plot(lambda_grid, func_gcv_abs_med_grid, ...
%     'DisplayName', 'GCV Median absolute error');
plot(lambda_grid, 0.5 * predintval_length_oos_grid, ...
    'LineWidth', 1.5, ...
    'Color', [0.4940 0.1840 0.5560], ...
    'DisplayName', 'Prediction (half) interval length');
[~, idx] = min(predintval_length_oos_grid);
plot(lambda_grid(idx), 0.5 * predintval_length_oos_grid(idx), '*', ...
    'LineWidth', 1.5, ...
    'Color', [0.4940 0.1840 0.5560], 'HandleVisibility', 'off');
plot(lambda_grid, 0.5 * predintval_length_gcv_grid, ...
    'LineWidth', 1.5, ...
    'LineStyle', '-.', ...
    'Color', [0.4940 0.1840 0.5560], ...
    'HandleVisibility', 'off');
%     'DisplayName', 'GCV interval length');
[~, idx] = min(predintval_length_gcv_grid);
plot(lambda_grid(idx), 0.5 * predintval_length_gcv_grid(idx), 'o', ...
    'LineWidth', 1.5, ...
    'Color', [0.4940 0.1840 0.5560], 'HandleVisibility', 'off');
legend show
legend box off
legend('Location', 'Best')

xlabel('\lambda')
ylabel('Error')


%% save plot
save_str = strcat(...
    'tuning_vs_lambda', ...
    '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/tuning/', save_str);
saveas(gcf, save_fig_str)
saveas(gcf, save_fig_str, 'epsc')
saveas(gcf, save_fig_str, 'png')