function [sigma, h] = SteinPost(Data, h)
% function [sigma, h] = SteinPost(Data)
% function [sigma, h] = SteinPost(Data, h)
% 
% Fits the Stein-Optimal Bayesian Covariance Matrix estimator described in
% Gillen (2012) "Bayesian Minimization of Stein Loss in Covariance Matrix
% Estimation"
%
%Inputs:  
% Data (T x N)      - Data for T realizations of N random variables between
%                     which you want to estimate the covariances 
%
% [h] (1 x 1)       - Bandwidth for estimator.  If not provided, bandwidth
%                     will be calculated using simulated cross-validation 
%                     with min(10*N, 1000) simulated samples on a grid of 
%                     100 equally-spaced points between 0 and the mean
%                     variance in the sample.
%
%Outputs:
% sigma (N x N)  - Stein-Optimal Posterior Covariance Matrix Expectation
% h (1 x 1)      - Stein-optimal posterior bandwidth

[T, N] = size(Data);
[BetaHat, ~, SigmaF] = princomp(Data); 

if nargin == 1 %Fit Bandwidth for h
    M = min(10*N, 1000);
    K = 100;
    
    hLoss = zeros(M, K);
    
    TestSigma = cov(Data);
    hMax = mean(diag(TestSigma));
    hMin = hMax/K;
    BigH = hMin:((hMax - hMin)/K):hMax;
    
    for m = 1:M
        mData = mvnrnd(zeros(1, N), TestSigma, T);
        mMeanData = mean(mData);
        mData = mData - mMeanData(ones(T, 1), :);

        mVar = var(mData);
        [mBetaHat, ~, mSigmaF] = princomp(mData);

        for hh = 1:K
            hmse = BigH(hh);

            mhdelta = T*(mSigmaF.^2)./ ...
                    (T*(mSigmaF.^2) + N*(2*N+3)*(hmse^2)*(mSigmaF) + ...
                        N*(N-1)*(hmse^4)*ones(N, 1)); 

            mhsigma = mBetaHat*diag(mhdelta.*mSigmaF)*mBetaHat';
            mhsigma(logical(eye(N))) = mVar;

            hLoss(m, hh) = sum(sum((mhsigma - TestSigma).^2));
        end
    end

    meanhLoss = mean(hLoss);
    h = BigH(meanhLoss == min(meanhLoss));
end

delta = T*(SigmaF.^2)./ ...
        (T*(SigmaF.^2) + N*(2*N+3)*SigmaF*(h^2) + N*(N-1)*ones(N, 1)*(h^4)); 
sigma = BetaHat*diag( delta.*SigmaF )*BetaHat';
sigma(logical(eye(N))) = (var(Data))';

end