clear variables

% number of samples and features
n = 1000;
gamma = 2;
p = gamma * n;

% ar-1 covariance
rho = 0.25;
Sigma = toeplitz(rho.^(0:1:p-1));
% Sigma = Sigma;
[W, R] = eig(Sigma);
% reorder the eigenvalues and eigenvectors of Sigma
[~, ind] = sort(diag(R), 'descend');
R = R(ind, ind);
W = W(:, ind);
dR = diag(R);

% ar-1 features
Z = randn(n, p);
% X = Z * sqrtm(Sigma);
% X = Z * W * diag(sqrt(dR)) * W';
% skip Z * W as it has the same distribution as Z
tZ = Z * W;
X = tZ * (sqrt(dR) .* W');
[U, sqrtS, V] = svd(X/sqrt(n));
% S = sqrtS' * sqrtS;
% get the first n entries of S
% dS = diag(S);
% dSn = dS(1:n);
dSn = diag(sqrtS).^2;

sig_energy = 1;
noi_energy = 0;

% signal and response
beta0 = W(:, 1) / sqrt(dR(1)) * sqrt(sig_energy);
% beta0 = W(:, p) / sqrt(dR(p)) * sqrt(sig_energy);
% beta0 = W(:, p) / sqrt(R(p, p));
% beta0 = W(:, n) / sqrt(dR(n)) * sqrt(sig_energy);
% beta0 = randn(p, 1) * sqrt(sig_energy) / sqrt(p);

y = zeros(n, 1);

for i = 1:n
    y(i) = (beta0' * X(i, :)')^2 - sig_energy;
end

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

% 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 = zeros(n_test, 1);

for j = 1:n_test
    y_test(j) = (beta0' * X_test(j, :)')^2 - sig_energy;
end

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(resid_gcv .^2 );
    func_loo = mean(resid_loo .^2);

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

    hy_test = X_test * ridge_estim;

    resid_oos = y_test - hy_test;

    func_oos = mean(resid_oos .^2 );
    
    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('MSE')

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

% figure
% hold on

