clear all

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

% ar-1 covariance
rho = 0.48;
% Sigma = toeplitz(rho.^(0:1:p-1));
Sigma = (1 - rho) * eye(p) + rho * ones(p, 1) * ones(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;

hSigma = X' * X / n;

n_grid = 100;

lambda_min = -1 * min(dSn) + 0.01;
lambda_max = 2;
lambda_grid = linspace(lambda_min, lambda_max, n_grid);

pop_functional = zeros(n_grid, 1);
emp_functional = zeros(n_grid, 1);

for i = 1:length(lambda_grid)
    
    pop_functional(i) = 1 / (1 + trace(Sigma * pinv(hSigma + lambda_grid(i) * eye(p))) / n);
    emp_functional(i) = 1 - trace(hSigma * pinv(hSigma + lambda_grid(i) * eye(p))) / n ;
end

plot(lambda_grid, pop_functional)
hold on
plot(lambda_grid, emp_functional)
xline(lambda_min)

legend('pop', 'emp', 'min eig')

