function [VD,VDu,VDd] = var_vd(Ainv,B,Bs,Vs,Om,h,ci,seq,tr)
% Computes variance decomposition for structural VAR(p)
%   Input:  Ainv  = structural restriction matrix
%           B     = OLS estimates
%           Bs    = bs OLS estimates (enter scalar for no bs option)
%           Vs    = bs covariance estimates
%           Om    = cov matrix for structural shocks
%           h     = forecast horizon (default = 20)
%           ci    = confidence level (default = 0.95)
%           seq   = compute VD at all horizons, 1 = yes, 0 = no (optional)
%           tr    = type of trend: const,linear,quadratic = (0,1,2)
%   Output: VD    = k-by-k matrix of decompositions
%           VDu   = upper confidence band (0 if no bs option)
%           VDd   = lower confidence band (0 if no bs option)
%   Author: Kyle Jurado (kej2108@columbia.edu)

%% Options
if nargin == 5; h  = 20; ci = 0.95; seq = 0; tr = 0; end
if nargin == 6; ci = 0.95; seq = 0; tr = 0; end
if nargin == 7; seq = 0; tr = 0; end
if nargin == 8; tr = 0; end

%% Initialization
k   = size(B,2);                                    %number of variables
p   = (length(B)-(tr+1))/k;                         %number of lags
J   = [eye(k),zeros(k,(p-1)*k)];                    %selection matrix
I  = [eye(k*(p-1)),zeros(k*(p-1),k)];               %block diagonal
C  = [B(tr+2:end,:)';I];                            %companion matrix
DOm = cell(h,1);                                    %rf MA coef matrix
Q   = cell(k,1);                                    %selection matrix
for i = 1:k
    Q{i}      = zeros(k,k);      
    Q{i}(i,i) = 1;
end
%% Variance decomposition (all horizons)
if seq == 1;
VD = zeros(k,k,h);                                  %initialize arrays
vd = zeros(k,1,k);
for fh = 1:h                                        %outer loop
    M   = 0;                                        %initialize dummy
    for t = 1:fh                                    %total fe variance loop
        DOm{t} = (J*C^(t-1)*J')*Ainv*Om;            %MA coefs times sdevs
        M      = M + DOm{t}*DOm{t}';                %add all multipliers
    end
    MSE   = diag(M);                                %keep variances only
    for i = 1:k                                     %for each variable
        M  = 0;                                     %initialize dummy
        for t = 1:fh                                %contribution loop
            M = M+DOm{t}*Q{i}*DOm{t}';              %part from variable i
        end
        vd(:,:,i) = diag(M)./MSE;                   %fraction contributed
    end
    VD(:,:,fh) = squeeze(vd);                       %compress to matrix
end

%% Confidence Bands (optional)
if (1-length(Bs))==0; VDu = 0; VDd = 0; end;
if (1-length(Bs))~=0;
S   = length(Bs);                                   %number of bs samples
VDb = zeros(k,k,h,S);                               %initialize array
vdb = zeros(k,1,k);
for s = 1:S
    tic;
    Cb = [Bs{s}(2:end,:)';I];                       %companion matrix
    Pb    = chol(Vs{s},'lower');
    Omb   = diag(diag(Pb));
    Ainvb = tril(Pb/Omb);    
for fh = 1:h                                        %outer loop
    M   = 0;                                        %initialize dummy
    for t = 1:fh                                    %total fe variance loop
        DOm{t} = (J*Cb^(t-1)*J')*Ainvb*Omb;         %MA coefs times sdevs
        M      = M + DOm{t}*DOm{t}';                %add all multipliers
    end
    MSE   = diag(M);                                %keep variances only
    for i = 1:k                                     %for each variable
        M  = 0;                                     %initialize dummy
        for t = 1:fh                                %contribution loop
            M = M+DOm{t}*Q{i}*DOm{t}';              %part from variable i
        end
        vdb(:,:,i) = diag(M)./MSE;                  %fraction contributed
    end
    VDb(:,:,fh,s) = squeeze(vdb);                   %compress to matrix
end
VDs = sort(VDb,4);
VDu = VDs(:,:,:,ceil(S*(1+ci)/2));                  %upper band
VDd = VDs(:,:,:,ceil(S*(1-ci)/2));                  %lower ('down') band
if s  == 1; disp('VD bootstrap simulations'); end;
fprintf('Simulation %d, Elapsed Time %f \n',s,toc);
end
end
end

%% Variance decompositions (one horizon only)
if seq == 0;
VD  = zeros(k,k);                                   %initialize matrix
M   = 0;                                            %initialize dummy
for t = 1:h
    DOm{t} = (J*C^(t-1)*J')*Ainv*Om;                %MA coefs times sdevs
    M      = M + DOm{t}*DOm{t}';                    %add all multipliers
end
MSE   = diag(M);                                    %keep variances only
for i = 1:k
    M  = 0;                                         %initialize dummy
    for t = 1:h
        M = M+DOm{t}*Q{i}*DOm{t}';                  %part from variable i
    end
    VD(:,i) = diag(M)./MSE;                         %fraction contributed
end

%% Confidence bands (optional)
if (1-length(Bs))==0; VDu = 0; VDd = 0; end;
if (1-length(Bs))~=0;
S   = length(Bs);                                   %number of bs samples
VDb = zeros(k,k,S);                                 %initialize array
for s = 1:S
    Cb  = [Bs{s}(2:end,:)';I];                      %companion matrix
    Pb    = chol(Vs{s},'lower');
    Omb   = diag(diag(Pb));
    Ainvb = tril(Pb/Omb);    
    M   = 0;                                        %initialize dummy
    for t = 1:h
        DOm{t} = (J*Cb^(t-1)*J')*Ainvb*Omb;         %MA coefs times sdevs
        M      = M + DOm{t}*DOm{t}';                %add all multipliers
    end
    MSE   = diag(M);                                %keep variances only
    for i = 1:k
        M  = 0;                                     %initialize dummy
        for t = 1:h
            M = M+DOm{t}*Q{i}*DOm{t}';              %part from variable i
        end
        VDb(:,i,s) = diag(M)./MSE;                  %fraction contributed
    end
if s  == 1; disp('VD bootstrap simulations'); end;
fprintf('Simulation %d, Elapsed Time %f \n',s,toc);
end
VDs = sort(VDb,3);
VDu = VDs(:,:,ceil(S*(1+ci)/2));                    %upper band
VDd = VDs(:,:,ceil(S*(1-ci)/2));                    %lower ('down') band
end
end
end