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 = 100;
noi_energy = 0.5;

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

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


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;

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

figure
hold on
% h_insample = cdfplot(y_new - hy);
% h_insample.LineWidth = 2;
% h_insample.DisplayName = 'in-sample (same X, new y)';
% h_orig = cdfplot(resid_orig);
% h_orig.LineWidth = 2;
% h_orig.DisplayName = 'fitted uncorrected';
% h_orig = cdfplot(resid_orig .* (1 + noi_energy * diag(smooth_mat)));
% h_orig.LineWidth = 2;
% h_orig.DisplayName = 'Mallow-type corrected';
h_test = cdfplot((y_test - hy_test) / norm(y_test - hy_test) * sqrt(n_test));
h_test.LineWidth = 1;
% h_test.Marker = '*';
h_test.DisplayName = 'out-of-sample (new X, y)';
h_loo = cdfplot((resid_loo) / norm(resid_loo) * sqrt(length(resid_loo)));
h_loo.LineWidth = 1;
h_loo.DisplayName = 'loo corrected';
h_gcv = cdfplot((resid_gcv) / norm(resid_gcv) * sqrt(length(resid_gcv)));
h_gcv.LineWidth = 1;
h_gcv.DisplayName = 'gcv corrected';

title('cumulative distributions of residuals')

xlim([-5, 5])

x_grid = linspace(-10, 10, 100);
plot(x_grid, normcdf(x_grid), 'LineWidth', 1, 'DisplayName', 'Normal CDF')

legend('show', 'Location', 'best')