%% clear all variables and start fresh
clear variables

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

%% problem parameters
n = 1000;
p = 2000;

n_test = 1e5;

lambda_min = 0.01;
lambda_max = 1;
n_grid = 50;

lambda_grid = linspace(lambda_min, lambda_max, n_grid);

%% generate train and test features

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

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 = 5;
beta0 = sqrt(sig_energy) * generate_signal_vector(p, sig_type);

% 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 = 'nlsq';
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);


%% vectors to store variour quantities

quantupper_gcv = zeros(length(lambda_grid), 1);
quantmiddle_gcv = zeros(length(lambda_grid), 1);
quantlower_gcv = zeros(length(lambda_grid), 1);
mean_gcv = zeros(length(lambda_grid), 1);

quantupper_oos = zeros(length(lambda_grid), 1);
quantmiddle_oos = zeros(length(lambda_grid), 1);
quantlower_oos = zeros(length(lambda_grid), 1);
mean_oos = zeros(length(lambda_grid), 1);

coverage = 0.5;
misscoverge = 1 - coverage;

for i = 1:n_grid
    
    disp(i)
    
    lambda = lambda_grid(i);
    
    
    %% fit estimator and get smoother

    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_test = 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)
        
            gram_mat = X * X' / n;
            resid_mod = gram_mat \ y;

            resid_gcv = resid_mod / (trace(pinv(gram_mat)) / n);
    %         resid_loo = resid_mod ./ diag(pinv(gram_mat));
    else
        
            ridge_smoother = get_ridge_smoother(X, lambda);
            hy = predict_linestim(ridge_estim, X);
            resid_orig = y - hy;
            resid_gcv = resid_orig / (1 - trace(ridge_smoother) / n);
    end
    
    resid_gcv_sq = resid_gcv .^2;
    resid_test_sq = resid_test .^2;
    
    quantupper_gcv(i, 1) = quantile(resid_gcv_sq, 1 - misscoverge/2);
    quantmiddle_gcv(i, 1) = quantile(resid_gcv_sq, 0.5);
    quantlower_gcv(i, 1) = quantile(resid_gcv_sq, misscoverge/2);
    mean_gcv(i, 1) = mean(resid_gcv .^2);
    
    quantupper_oos(i, 1) = quantile(resid_test_sq, 1 - misscoverge/2);
    quantmiddle_oos(i, 1) = quantile(resid_test_sq, 0.5);
    quantlower_oos(i, 1) = quantile(resid_test_sq, misscoverge/2);
    mean_oos(i, 1) = mean(resid_test .^2);

end

%% plot different curves

plot(lambda_grid, mean_oos, ...
    'Color', [0 0.4470 0.7410], ...
    'DisplayName', 'Test mean')
hold on
grid on
% semilogy(gamma_grid, quantmiddle_oos, ...
%     'DisplayName', 'Test median')
plot(lambda_grid, quantupper_oos, ...
    'Color', [0.8500 0.3250 0.0980], ...
    'DisplayName', 'Test upper quantile')
plot(lambda_grid, quantlower_oos, ...
    'Color', [0.9290 0.6940 0.1250], ...
    'DisplayName', 'Test lower quantile')
plot(lambda_grid, mean_gcv, '-*', ...
    'Color', [0 0.4470 0.7410], ...
    'DisplayName', 'GCV mean')
% semilogy(gamma_grid, quantmiddle_gcv, '-*', ...
%     'DisplayName', 'GCV median')
plot(lambda_grid, quantupper_gcv, '-*', ...
    'Color', [0.8500 0.3250 0.0980], ...
    'DisplayName', 'GCV upper quantile')
plot(lambda_grid, quantlower_gcv, '-*', ...
    'Color', [0.9290 0.6940 0.1250], ...
    'DisplayName', 'GCV lower quantile')

legend show
legend box off
legend('Location', 'Best')
xlabel('\lambda')
ylabel('Squared error')


save_str = strcat(...
    'predintval_vs_lambda', ...
    'n', num2str(n), '_', ...
    'nt', num2str(n_test), '_', ...
    'featcov', feat_cov_type, '_', ...
    'feat', feat_type, '_', ...
    'featparam', num2str(feat_param), '_', ...
    'sig', sig_type, '_', ...
    'noi', noi_type, '_', ...
    'noiparam', num2str(noi_param), '_', ...
    'mod', mod_type, '_', ...
    'lam', num2str(lambda), '_', ...
    'seed', num2str(seed));
% 
% save_folder = '../figures/';
% 
% saveas(h_oos, fullfile(save_folder, save_str), 'epsc')

saveas(gcf, save_str, 'epsc')
saveas(gcf, save_str, 'png')