%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%% Demo for shrinkage covariance estimators   %%%
%%%             Written by Yilun Chen          %%%
%%%                yilun@umich.edu             %%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



clear
clc
close all

err_LW_arr = zeros(200,1);
err_S_arr = zeros(200,1);
err_new_arr = zeros(200,1);
err_arr = zeros(200,1);
err_true_arr = zeros(200,1);

r_err_arr = zeros(200,1);

rho_LW_arr = zeros(200,1);
rho_new_arr = zeros(200,1);
rho_arr = zeros(200,1);
rho_true_arr = zeros(200,1);

var_LW_arr = zeros(200,1);
var_new_arr = zeros(200,1);
var_arr = zeros(200,1);

seq = [5,7,10:10:120];
N_max = 100;
M_max = 50;
THat = zeros(length(seq),N_max);
S2Hat = zeros(length(seq),N_max);
STHat = zeros(length(seq),N_max);
rho_a_Hat = zeros(length(seq),N_max);
rho_b_Hat = zeros(length(seq),N_max);
S1Hat = zeros(length(seq),N_max);
TaoHat = zeros(length(seq),N_max);

rho_I = zeros(1,max(seq));
rho_G = zeros(1,max(seq));


H = 0.8 ;
for n = seq
    n
    p = 100;
    %     n = 10;
    A = zeros(p,p);
    for ii = 1:p
        for jj = 1:p
            A(ii,jj) = 0.5^(abs(ii-jj));
        end
    end
    
    %%%or%%%%
%     A = eye(p);
%     for ii = 1:p
%         for jj = ii+1:p
%             A(ii,jj) = 1/2*((abs(ii-jj)+1)^(2*H)-2*abs(ii-jj)^(2*H)+(abs(ii-jj)-1)^(2*H));
%             A(jj,ii) = A(ii,jj);
%         end
%     end
    
    
    SIGMA = A;
    L = chol(SIGMA);
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    err_LW = 0;
    err_S = 0;
    err_new = 0;
    err_ = 0;
    err_true = 0;
    r_err = 0;
    
    C_np = n/(n+2);
    
    rho_I(n) = ((n+3)/n/(n+1)*trace(SIGMA*SIGMA)+1/n*trace(SIGMA)^2)/...
        ((1+1/n-2/n/p)*trace(SIGMA*SIGMA)+(1/n-1/p)*trace(SIGMA)^2);
    rho_G(n) = (n-1)/n^2*(trace(SIGMA*SIGMA)+trace(SIGMA)^2)/...
        ((1+1/n-2/n/p)*trace(SIGMA*SIGMA)+(1/n-1/p)*trace(SIGMA)^2);
    
    for iii = 1:N_max
        X = randn(p,n);
        X = L'*X;
        S = X*X'/n;
        mu = trace(S)/p;
        rho = 0;

        for ii = 1:n
            rho = rho+norm(X(:,ii)*X(:,ii)'-S,'fro')^2;
        end
        rho = rho/n^2/(norm(S-mu*eye(p),'fro')^2);
        rho = min(rho,1);

        R_LW = (1-rho)*S+rho*mu*eye(p);

        err_LW = err_LW + norm(R_LW-SIGMA,'fro')^2;
        err_S = err_S+norm(S-SIGMA,'fro')^2;

        rho_new = 0;
        for ii = 1:n
            rho_new = rho_new+C_np*(2*trace(S*S)+trace(S)^2)-trace(S*S);
        end
        rho_new = rho_new/n^2/(norm(S-mu*eye(p),'fro')^2);

        rho_new = min(rho_new,1);
        rho_new = max(rho_new,0);
        R_new = (1-rho_new)*S+rho_new*mu*eye(p);

        rho_ = (1/(n+2)*trace(S)^2+(n-2)/n/(n+2)*trace(S*S))/(trace(S*S)-1/p*trace(S)^2);

        rho_ = rho_new;
        rho_ = (1/(n+1)*trace(S)^2+1/(n+1)*trace(S*S))/(trace(S*S)-1/p*trace(S)^2);

        rho_  = min(rho_,1);
        SIGMA_hat = (1-rho_)*S + rho_*mu*eye(p);
        R_ = SIGMA_hat;

        %%%%%
        rho_true = (trace(SIGMA)^2+trace(SIGMA*SIGMA))/(n*(trace(SIGMA*SIGMA)-1/p*trace(SIGMA)^2)+trace(SIGMA)^2+trace(SIGMA*SIGMA));
        R_true = (1-rho_true)*S+rho_true*trace(S)/p*eye(p);
        %%%%%
        err_ = err_ + norm(R_-SIGMA,'fro')^2;
        err_new = err_new+norm(R_new-SIGMA,'fro')^2;
        err_true = err_true + norm(R_true-SIGMA,'fro')^2;

        r_err = r_err + norm(R_true-R_,'fro')^2;

        rho_LW_arr(n) = rho_LW_arr(n) + rho;
        rho_arr(n) = rho_arr(n) + rho_;
        rho_true_arr(n) = rho_true_arr(n) + rho_true;
        rho_new_arr(n) = rho_new_arr(n) + rho_new;

        var_LW_arr(n) = var_LW_arr(n) + rho^2;
        var_arr(n) = var_arr(n) + rho_^2;
        var_new_arr(n) = var_new_arr(n) + rho_new^2;
        THat(n,iii) = 1/(p-1)*(p*trace(S^2)/trace(S)^2-1);
        S1Hat(n,iii) = trace(S*S);
        S2Hat(n,iii) = trace(S)^2;
        STHat(n,iii) = S2Hat(n,iii)/THat(n,iii);
        rho_a_Hat(n,iii) = rho_;
        rho_b_Hat(n,iii) = rho_new;
    end
    disp('true')
    err_true/N_max
    disp('_')
    err_/N_max
    disp('new')
    err_new/N_max
    disp('honey')
    err_LW/N_max
    disp('sample')
    err_S/N_max

    err_true_arr(n) = err_true/N_max;
    err_arr(n) = err_/N_max;
    err_new_arr(n) = err_new/N_max;
    err_LW_arr(n) = err_LW/N_max;
    err_S_arr(n) = err_S/N_max;
    r_err_arr(n) = r_err/N_max;

    rho_LW_arr(n) = rho_LW_arr(n)/N_max;
    rho_arr(n) = rho_arr(n)/N_max;
    rho_true_arr(n) = rho_true_arr(n)/N_max;
    rho_new_arr(n) = rho_new_arr(n)/N_max;

    var_LW_arr(n) = var_LW_arr(n)/N_max - rho_LW_arr(n)^2;
    var_arr(n) = var_arr(n)/N_max - rho_arr(n)^2;
    var_new_arr(n) = var_new_arr(n)/N_max - rho_new_arr(n)^2;

end
err_true_arr_ = err_true_arr(seq);
err_new_arr_ = err_new_arr(seq);
err_arr_ = err_arr(seq);
err_LW_arr_ = err_LW_arr(seq);
err_S_arr_ = err_S_arr(seq);

figure,plot(seq,err_true_arr_,'*-');
hold
plot(seq,err_arr_,'ro-')

% plot(seq,err_S_arr_,'rd')
plot(seq,err_new_arr_,'b^-')
plot(seq,err_LW_arr_,'sg-')
legend('Oracle','OAS','RBLW','LW')
grid on
xlabel('n')
ylabel('MSE')

rho_LW_arr_ = rho_LW_arr(seq);
rho_arr_ = rho_arr(seq);
rho_true_arr_ = rho_true_arr(seq);
rho_new_arr_ = rho_new_arr(seq);

figure,plot(seq,rho_true_arr_,'*-');
hold
plot(seq,rho_arr_,'ro-')

plot(seq,rho_new_arr_,'b^-')
plot(seq,rho_LW_arr_,'sg-')

legend('Oracle','OAS','RBLW','LW');
grid on
xlabel('n')
ylabel('Shrinkage coefficient')
