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));
% equicorrelated covariance
rho = 0.1;
Sigma = (1 - rho) * eye(p) + rho * ones(p, 1) * ones(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 = 0.1;
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);
eps = sqrt(noi_energy) * binornd(1, 0.5, [n 1]);
y_sig = X * beta0;

% nu = 3; % degrees of freedom for the t distribution
% eps = sqrt(noi_energy) / sqrt(nu / (nu - 2)) * trnd(nu, [n 1]);

% stable distribution with slight alpha and beta
% stable_pd = makedist('Stable', 'alpha', 1.5, 'beta', 0.5);

% eps = random(stable_pd, [n 1]);
% % histogram(eps, 'Normalization', 'pdf')

% y = X * beta0 + eps;
y = y_sig + eps;

% lambda = 1;

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

% lambda_grid = linspace(lambda_min, lambda_max, n_grid);

% lambda_grid = [lambda_grid, logspace(0, 3, 10)];

lambda_grid = [1];

func_gcv_abs_grid = zeros(n_grid, 1);
func_loo_abs_grid = zeros(n_grid, 1);
func_oos_abs_grid = zeros(n_grid, 1);

% median of the absolute values
func_gcv_abs_med_grid = zeros(n_grid, 1);
func_loo_abs_med_grid = zeros(n_grid, 1);
func_oos_abs_med_grid = zeros(n_grid, 1);

func_gcv_sq_grid = zeros(n_grid, 1);
func_loo_sq_grid = zeros(n_grid, 1);
func_oos_sq_grid = zeros(n_grid, 1);

n_test = 100000;

% % new response at the same design points
% 
% %     eps_new = noi_energy * randn(n, 1);
% eps_new = sqrt(noi_energy) / sqrt(nu / (nu - 2)) * trnd(nu, [n 1]);
% y_new = X * beta0 + eps_new;

% generate new points

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

X_test = Z_test * sqrtm(Sigma);
% eps_test = sqrt(noi_energy) * randn(n_test, 1);
eps_test = sqrt(noi_energy) * binornd(1, 0.5, [n_test 1]);
% eps_test = sqrt(noi_energy) / sqrt(nu / (nu - 2)) * trnd(nu, [n_test 1]);
% eps_test = random(stable_pd, [n_test 1]);
y_test = X_test * beta0 + eps_test;

% histogram(eps_test, 'Normalization', 'pdf')

% matrices to store residuals
resid_orig_grid = zeros(n, n_grid);
resid_gcv_grid = zeros(n, n_grid);
resid_loo_grid = zeros(n, n_grid);
resid_oos_grid = zeros(n_test, n_grid);

for i = 1:length(lambda_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));

    resid_orig_grid(:, i) = resid_orig;
    resid_gcv_grid(:, i) = resid_gcv;
    resid_loo_grid(:, i) = resid_loo;
    
    func_gcv_abs = mean(abs(resid_gcv));
    func_loo_abs = mean(abs(resid_loo));

    func_gcv_abs_grid(i) = func_gcv_abs;
    func_loo_abs_grid(i) = func_loo_abs;
    
    func_gcv_abs_med = median(abs(resid_gcv));
    func_loo_abs_med = median(abs(resid_loo));
    
    func_gcv_abs_med_grid(i) = func_gcv_abs_med;
    func_loo_abs_med_grid(i) = func_loo_abs_med;
    
    func_gcv_sq = mean(resid_gcv.^2);
    func_loo_sq = mean(resid_loo.^2);

    func_gcv_sq_grid(i) = func_gcv_sq;
    func_loo_sq_grid(i) = func_loo_sq;
    
    inner_mat = [diag(diag(sqrtS) ./ (dSn + lambda)); zeros(p-n, n)];
    ridge_estim = V * inner_mat * U' * y / sqrt(n);

    hy_test = X_test * ridge_estim;

    resid_oos = y_test - hy_test;

    resid_oos_grid(:, i) = resid_oos;
    
    func_oos_abs = mean(abs(resid_oos));
    func_oos_abs_grid(i) = func_oos_abs;
    
    func_oos_abs_med = median(abs(resid_oos));
    func_oos_abs_med_grid(i) = func_oos_abs_med;
    
    func_oos_sq = mean(resid_oos.^2);
    func_oos_sq_grid(i) = func_oos_sq;
    

end

% figure
% grid on
% hold on
% 
% plot(lambda_grid, func_oos_abs_grid, 'DisplayName', 'oos abs');
% plot(lambda_grid, func_loo_abs_grid, 'DisplayName', 'loo abs');
% plot(lambda_grid, func_gcv_abs_grid, 'DisplayName', 'gcv abs');
% 
% plot(lambda_grid, func_oos_abs_med_grid, 'DisplayName', 'oos med')
% plot(lambda_grid, func_gcv_abs_med_grid, 'DisplayName', 'gcv med')
% plot(lambda_grid, func_loo_abs_med_grid, 'DisplayName', 'loo med')
% 
% plot(lambda_grid, sqrt(func_oos_sq_grid), 'DisplayName', 'oos sqrt');
% plot(lambda_grid, sqrt(func_loo_sq_grid), 'DisplayName', 'loo sqrt');
% plot(lambda_grid, sqrt(func_gcv_sq_grid), 'DisplayName', 'gcv sqrt');
% 
% legend show
% 
% xlabel('\lambda')
% ylabel('absolute value error')
% 
% hold off

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

% figure
% hold on

edges_fine = [-5 -5:0.1:5 5];
edges_coarse = [-5 -5:0.25: 5];
figure
h_oos = histogram(resid_oos_grid(:,1), edges_fine, 'Normalization', 'pdf');
hold on
h_gcv = histogram(resid_gcv_grid(:,1), edges_coarse, 'Normalization', 'pdf');
h_loo = histogram(resid_loo_grid(:,1), edges_coarse, 'Normalization', 'pdf');
h_orig = histogram(resid_orig(:,1), edges_coarse, 'Normalization', 'pdf');
legend('oos erros', 'gcv errors', 'loo errors', 'raw errors')
h_gcv.BinWidth = 0.25;
h_loo.BinWidth = h_gcv.BinWidth;
h_orig.BinWidth = h_gcv.BinWidth;
% ylim([-Inf, 0.5])
grid on

