% number of samples
n = 500;

% ratio of p/n
gamma = 2;

% number of features
p = gamma * n;

% ar-1 covariance
rho = 0.1;
r = rho.^(0:1:p-1);
Sigma = toeplitz(r);
[V_Sigma, D_Sigma] = eig(Sigma);

% ar-1 features
Z = randn(n, p);
X = Z * sqrtm(Sigma);

% Sigma = eye(p);
% X = Z * Sigma;

% sample covariance
hSigma = X' * X / n;
[V_hSigma, D_hSigma] = eig(hSigma);

% error vector
% err = randn(n, 1);

err = zeros(n, 1);

% signal aligned with top eigenvector
beta0 = V_Sigma(:, p) / sqrt(p);

% Sigma_beta0 = Sigma;
% beta0 = V_Sigma * diag(D_Sigma) / sqrt(p);
% beta0 = sqrtm(Sigma_beta0) * randn(p, 1) / sqrt(p);
% beta0 = randn(p, 1) / sqrt(p);
 
% Sigma_beta0 = V_Sigma(:, p) * V_Sigma(:, p)';

% Sigma_beta0 = V_Sigma * diag(flip(diag(D_Sigma))) * V_Sigma';

% Sigma_beta0 = V_Sigma * D_Sigma * V_Sigma';

% response
y = X * beta0 + err;

% ridge estimator, preduon risk
lambda_grid = -50:1:50;
pred_risk = zeros(length(lambda_grid), 1);
gcv = zeros(length(lambda_grid), 1);

for i = 1:length(lambda_grid)
   
    lambda = lambda_grid(i);
    hSigma_lambda = (hSigma + lambda * eye(p));
%     hbeta = (hSigma + lambda * eye(p)) \ (X' * y) / n;
    hbeta = pinv(hSigma_lambda) * X' * y / n;
    pred_risk(i) = (hbeta - beta0)' * Sigma * (hbeta - beta0);
    
%     hSigma_lambda = (hSigma + lambda * eye(p));
%     pred_risk(i) = ...
%     lambda^2 * trace(Sigma_beta0 * pinv(hSigma_lambda) * Sigma * pinv(hSigma_lambda)) / p;
%     
%     K_lambda = X * X'/n + lambda * eye(n);
%     gcv_num = trace(Sigma_beta0 * (K_lambda \ X)' * (K_lambda \ X)) / n / p;
%     gcv_den = (trace(inv(K_lambda)) / n)^2;
%     gcv(i) = gcv_num / gcv_den;
    
end

[~, idx] = min(pred_risk);
semilogy(lambda_grid, pred_risk, '-o')
xline(lambda_grid(idx), '-o')
grid on
hold on
plot(lambda_grid, gcv, '-*')
[~, idx] = min(gcv);
xline(lambda_grid(idx), '-*')
legend('risk', 'min risk', 'gcv', 'min gcv')
xlabel('\lambda')