%% 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 = 'iso';
% feat_cov_param = 0.25;
feat_cov = generate_feature_covariance(p, feat_cov_type);

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 = 2;
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 = 'iso';
% sig_energy = 50;
% beta0 = W(:, 1) / sqrt(dR(1)) * sqrt(sig_energy);

% 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);

%% 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 * min(dSn) + 0.01;
% lambda_min = -0.15;
% lambda_max = 1;
% n_lambda_grid = 31;
% 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

% 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);

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

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

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_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)
   
    resid_mod = gram_mat \ y;
    
    resid_gcv = resid_mod / (trace(pinv(gram_mat)) / n);
    resid_loo = resid_mod ./ diag(pinv(gram_mat));
    
end


%% form prediction intervals

% to store end points of the interval for gcv
predintval_gcv_grid = zeros(n_test, 2, length(coverage_grid));

% to store end points of the interval for loo
predintval_loo_grid = zeros(n_test, 2, length(coverage_grid));

% to store upper and lower qunatiles for gcv and loo
quanthigh_gcv_grid = zeros(length(coverage_grid), 1);
quantlow_gcv_grid = zeros(length(coverage_grid), 1);
quanthigh_loo_grid = zeros(length(coverage_grid), 1);
quantlow_loo_grid = zeros(length(coverage_grid), 1);

% to store coverage checks for gcv and loo
coverageind_gcv_grid = zeros(n_test, length(coverage_grid));
coverageind_loo_grid = zeros(n_test, length(coverage_grid));

for i = 1:length(coverage_grid)
    
    disp(i)

    quanthigh_gcv_grid(i) = quantile(resid_gcv, 1 - (1 - coverage_grid(i))/2);
    quanthigh_loo_grid(i) = quantile(resid_loo, 1 - (1 - coverage_grid(i))/2);

    quantlow_gcv_grid(i) = quantile(resid_gcv, (1- coverage_grid(i))/2);
    quantlow_loo_grid(i) = quantile(resid_loo, (1- coverage_grid(i))/2);

    predintval_gcv_grid(:, 1, i) = hy_test + quantlow_gcv_grid(i);
    predintval_gcv_grid(:, 2, i) = hy_test + quanthigh_gcv_grid(i);
    
    predintval_loo_grid(:, 1, i) = hy_test + quantlow_loo_grid(i);
    predintval_loo_grid(:, 2, i) = hy_test + quanthigh_loo_grid(i);
    
    coverageind_gcv_grid(:, i) = bitand( ...
        (y_test <= predintval_gcv_grid(:, 2, i)), ...
        (y_test >= predintval_gcv_grid(:, 1, i)));
    
    coverageind_loo_grid(:, i) = bitand( ...
        (y_test <= predintval_loo_grid(:, 2, i)), ...
        (y_test >= predintval_loo_grid(:, 1, i)));

    coverageval_gcv_grid(i, j) = mean(coverageind_gcv_grid(:, i));
    coverageval_loo_grid(i, j) = mean(coverageind_loo_grid(:, i));
    
    predintval_length_gcv_grid(i, j) = ...
        quanthigh_gcv_grid(i) - quantlow_gcv_grid(i);
    predintval_length_loo_grid(i, j) = ...
        quanthigh_loo_grid(i) - quantlow_loo_grid(i);
end

end
%% plots
marker_list = ['*', 'o', '+', 'x'];

% interval coverage
figure
hold on
grid on
for j = 1:length(lambda_grid)
    marker = marker_list(j);
    legend_str = ['\lambda = ', num2str(lambda_grid(j))];
    plot(coverage_grid, coverageval_loo_grid(:, j), ...
    'Marker', marker, ...
    'LineWidth', 1.5, ...
    'DisplayName', legend_str);
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

legend box off
legend('Location', 'Best')
xlabel('Nominal coverage')
ylabel('Empirical coverage')

%% save plot
save_str = strcat(...
    'coverage_comp_loo_', ...
    '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(lambda_grid)
    marker = marker_list(j);
    legend_str = ['\lambda = ', num2str(lambda_grid(j))];
plot(coverage_grid, predintval_length_loo_grid(:,j), ...
    'Marker', marker, ...
    'LineWidth', 2, ...
    'DisplayName', legend_str);

% plot(coverage_grid, predintval_length_loo_grid, ...
%     '*-', ...
%     'DisplayName', 'loo');
end
legend show
legend box off
legend('Location', 'Best')
xlabel('Nominal coverage')
ylabel('Prediction interval length')

save_str = strcat(...
    'length_comp_loo_', ...
    '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')