clearvars

% number of samples and features
n = 500; % number of smaples
gamma = 2; % overparametrization ratio; gamma := p / n
p = gamma * n; % not rounded off at the moment

% generate population covariance
% ar-1 covariance
rho = 0.25;
Sigma = toeplitz(rho.^(0:1:p-1));
[W, R] = eig(Sigma); % 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

% generate features
% we generate "random matrix theory" features
Z = randn(n, p); % n times p matrix with indep N(0, 1) entries
% Z = trnd(3, n, p); % n times p matrix with indep N(0, 1) entries
X = Z * sqrtm(Sigma);
[U, sqrtS, V] = svd(X/sqrt(n)); % svd of X / sqrt(n) = U * sqrtS * V'
dSn = diag(sqrtS).^2; % n eigenvalues of \hSigma

% parameters for signal generation
sig_energy = 25;
noi_energy = 1;

% signal and response
beta0 = sig_energy / p * randn(p, 1); % random signal with energy r^2
eps = noi_energy * randn(n, 1);
y = X * beta0 + eps;

% lambda = 1;

lambda_min = -1 * min(dSn) + 0.01;
lambda_max = 2.5;
n_grid = 10;

lambda_grid = linspace(lambda_min, lambda_max, n_grid);
quanthigh_gcv_grid = zeros(n_grid, 1);
quanthigh_loo_grid = zeros(n_grid, 1);
quanthigh_oos_grid = zeros(n_grid, 1);
quantlow_gcv_grid = zeros(n_grid, 1);
quantlow_loo_grid = zeros(n_grid, 1);
quantlow_oos_grid = zeros(n_grid, 1);

for i = 1:n_grid
   
    lambda = lambda_grid(i);
    
    smooth_mat = U * diag(dSn ./ (dSn + lambda))  * U';

    hy = smooth_mat * y;

    resid_orig = y - hy;
    resid_gcv = resid_orig ./ (1 - trace(smooth_mat) / n);
    resid_loo = resid_orig ./ (1 - diag(smooth_mat));

    quanthigh_gcv_grid(i) = quantile(resid_gcv, 0.975);
    quanthigh_loo_grid(i) = quantile(resid_loo, 0.975);
    
    quantlow_gcv_grid(i) = quantile(resid_gcv, 0.025);
    quantlow_loo_grid(i) = quantile(resid_loo, 0.025);
    
    inner_mat = [diag(diag(sqrtS) ./ (dSn + lambda)); zeros(p-n, n)];
    ridge_estim = V * inner_mat * U' * y / sqrt(n);

    % new response at the same design points

    eps_new = noi_energy * randn(n, 1);
    y_new = X * beta0 + eps_new;

    % generate new points

    n_test = 100000;
    Z_test = randn(n_test, p);
    % Z_test = trnd(3, n_test, p);

    X_test = Z_test * sqrtm(Sigma);
    eps_test = noi_energy * randn(n_test, 1);
    y_test = X_test * beta0 + eps_test;

    hy_test = X_test * ridge_estim;

    resid_oos = y_test - hy_test;
    
    quanthigh_oos = quantile(resid_oos, 0.975);
    quantlow_oos = quantiel(resid_oos, 0.025);

end

figure
grid on
hold on

plot(lambda_grid, quanthigh_oos_grid, 'DisplayName', 'oos upper');
plot(lambda_grid, quanthigh_loo_grid, 'DisplayName', 'loo upper');
plot(lambda_grid, quanthigh_gcv_grid, 'DisplayName', 'gcv upper');

plot(lambda_grid, quantlow_oos_grid, 'DisplayName', 'oos lower');
plot(lambda_grid, quantlow_loo_grid, 'DisplayName', 'loo lower');
plot(lambda_grid, quantlow_gcv_grid, 'DisplayName', 'gcv lower');

legend show

xlabel('\lambda')
ylabel('quantile')
% histogram(y_test - hy_test, ...
%     'Normalization', 'cdf', ...
%     'DisplayStyle', 'stairs', ...
%     'LineWidth', 2, ...
%     'DisplayName', 'true out-of-sample')

% figure
% hold on
