% number of samples
n = 200;

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

% beta0 = V_Sigma * diag(D_Sigma) / sqrt(p);
beta0 = V_Sigma * sqrtm(D_Sigma) * randn(p, 1) / sqrt(p);
% beta0 = randn(p, 1) / sqrt(p);
 
% response
y = X * beta0 + err;

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

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

figure
plot(lambda_grid, pred_risk)
grid on
hold on
plot(lambda_grid, gcv)
legend('risk', 'gcv')
xlabel('\lambda')


bias_inner_risk = zeros(length(lambda_grid), 1);
bias_inner_gcv = zeros(length(lambda_grid), 1);

for i = 1:length(lambda_grid)
   
    lambda = lambda_grid(i);
    hSigma_lambda = (hSigma + lambda * eye(p));
    K_lambda = X * X'/n + lambda * eye(n);
    
    bias_inner_risk(i) = ...
        lambda^2 * trace(Sigma * pinv(hSigma_lambda) * Sigma * pinv(hSigma_lambda)) / p;
    
%     bias_inner_gcv_num = ...
%         trace((hSigma_lambda \ sqrtm(hSigma)) * (hSigma_lambda \ sqrtm(hSigma))') / p;
    
%     bias_inner_gcv_den = ...
%         (1 - trace(hSigma_lambda \ hSigma) / n)^2;
    
    bias_inner_gcv_num = ...
        trace(Sigma * (X' / K_lambda) * (K_lambda \ X)) / n / p;

    bias_inner_gcv_den = ...
        (trace(inv(K_lambda)) / n)^2;
        
    bias_inner_gcv(i) = bias_inner_gcv_num / bias_inner_gcv_den;
end

figure
plot(lambda_grid, bias_inner_risk)
hold on
grid on
plot(lambda_grid, bias_inner_gcv)
legend('risk', 'gcv')
xlabel('\lambda')

trace



