clearvars

% number of samples and features
n = 1000; % 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 = 60;
noi_energy = 1;

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

% lambda = 1;

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

lambda_grid = linspace(lambda_min, lambda_max, n_grid);
func_gcv_grid = zeros(n_grid, 1);
func_loo_grid = zeros(n_grid, 1);
func_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));

    func_gcv = mean(abs(resid_gcv));
    func_loo = mean(abs(resid_loo));

    func_gcv_grid(i) = func_gcv;
    func_loo_grid(i) = func_loo;
    
    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;

    func_oos = mean(abs(resid_oos));
    
    func_oos_grid(i) = func_oos;

end

figure
grid on
hold on

plot(lambda_grid, func_oos_grid, 'DisplayName', 'oos');
plot(lambda_grid, func_loo_grid, 'DisplayName', 'loo');
plot(lambda_grid, func_gcv_grid, 'DisplayName', 'gcv');

legend show

xlabel('\lambda')
ylabel('absolute value error')

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

% figure
% hold on
