%% EB estimation
function res = EB_est(x,P,optNaN,blocks)

r = P.r;
p = P.p;
max_iter = P.max_iter;
thresh = P.thresh;

OPTS.disp=0;

[xBal,indNaN] = remNaNs_spline(x,optNaN);

[T,N] = size(x);

res = EB_algo(xBal,r,blocks,max_iter,thresh);

resid = res.idio;
resid(indNaN) = NaN;
R = nanvar(resid);

C = [res.lambda zeros(N,r*(p-1))];

As = {};
Qs = {};
InitVs = {};

% common factor VAR estimation

z = res.facs{1};
Z = [];
for kk = 1:p
    Z = [Z z(p-kk+1:end-kk,:)]; % stacked regressors (lagged SPC)
end
z = z(p+1:end,:);
% run the var chi(t) = A*chi(t-1) + e(t);

A_temp = inv(Z'*Z)*Z'*z;
A_c(1:r,1:r*p) = A_temp';
A_c(r+1:end,1:r*(p-1)) = eye(1:r*(p-1));
As{1} = A_c;

Q_c = zeros(r*p,r*p);
e = z  - Z*A_temp; % VAR residuals
Q_c(1:r,1:r) = cov(e);  % VAR residuals covariance matrix
Qs{1} = Q_c;

rp2 = (r*p)^2;

% Initial conditions for the block b
initV_c = reshape(inv(eye(rp2)-kron(A_c,A_c))*Q_c(:),r*p,r*p);
initVs{1} = initV_c;


for b =1:size(blocks,1)
    facs_b = res.facs{b+1};
    r_b = blocks{b,3};
    p_b = blocks{b,4};
    
    A_b = zeros(r_b*p_b,r_b*p_b);
    
    z = facs_b;
    Z = [];
    for kk = 1:p
        Z = [Z z(p_b-kk+1:end-kk,:)]; % stacked regressors (lagged SPC)
    end
    z = z(p_b+1:end,:);
    % run the var chi(t) = A*chi(t-1) + e(t);
    
    A_temp = inv(Z'*Z)*Z'*z;
    A_b(1:r_b,1:r_b*p_b) = A_temp';
    A_b(r_b+1:end,1:r_b*(p_b-1)) = eye(1:r_b*(p_b-1));
    As{b+1} = A_b;
    
    Q_b = zeros(r_b*p_b,r_b*p_b);
    e = z  - Z*A_temp; % VAR residuals
    Q_b(1:r_b,1:r_b) = cov(e);  % VAR residuals covariance matrix
    Qs{b+1} = Q_b;
    
    rp2 = (r_b*p_b)^2;
    
    % Initial conditions for the block b
    initV_b = reshape(inv(eye(rp2)-kron(A_b,A_b))*Q_b(:),r_b*p_b,r_b*p_b);
    initVs{b+1} = initV_b;
end

A = blkdiag(As{:});
Q = blkdiag(Qs{:});

r_cum = r + cumsum([0; cell2mat(blocks(:,3))]); %define the column indices for the factor blocks
r_tot = r_cum(end);

% Initial conditions
initV = blkdiag(initVs{:});
initZ = ones(r_tot,1); % %[randn(1,r*(nlag+1))]';

res.F = Z;

end

function res = EB_algo(Y,r,blocks,max_iter,thresh)

[T,N] = size(Y);

% CCA estimation of the common factors
sigs = {}; % empty cell to store block covariances
facs = {}; % empty cell to store block factors

% PC analysis block by block
for b = 1:size(blocks,1)
    yb = Y(:,blocks{b,2}); % restricting to block observation
    [v d] = eigs( cov(yb),r+blocks{b,3}); % extracting PCs
    fac = yb*v; % static proection to extract factors
    facs{b} = fac; % store factors
end

% select the couple of block whose first r (number of components)
% components have maximal correlation

rho = 0;
for i = 1:size(blocks,1)
    for j = (i+1):size(blocks,1)
        
        % for each pair i,j calculate canonical correlation between
        % the set of extracted factors
        
        fac1 = facs{i}; %extracted factors from the first block
        fac2 = facs{j}; %extracted factors from the second block
        
        T = length(fac1);
        
        s1 = 1/(T-1)*(fac1'*fac1);
        s2 = 1/(T-1)*(fac2'*fac2);
        s12 = 1/(T-1)*(fac1'*fac2);
        s21 = 1/(T-1)*(fac2'*fac1);
        
        A = s12*inv(s2)*s21; % matrix used to calculated the generalized eig. problem
        
        [vv dd] = eig(A,s1);
        
        if rho < max(max(dd)) %check if the new CC is better than the last one
            ind = [i j]; %extract the couple of indeces of the new best CC
            rho = max(max(dd)); %set the new CC as the benchmark
            load = vv; %extract loadings to calculate global factors at the end
        end
        
    end
end

glob = facs{ind(1)}*load(:,1:r); %extract global factors with extracted CCA parameters

common = glob*(inv(glob'*glob)*glob'*Y); %common component defined by glob
idio = Y - common;  %purging out the global component
r_cum = r + cumsum([0; cell2mat(blocks(:,3))]); %define the column indices for the factor blocks
tot_F = zeros(T,r_cum(end)); % preallocate the factor matrix of dim (T,r + sum_b(r_b))
tot_F(:,1:r) = glob; %fill in the first r columns with the extracted global factors

% blockwise factor extraction by means of PCA

for b = 1:size(blocks,1)
    
    % storing in local variables block related indices
    ind_b = blocks{b,2};
    r_cum_b = r_cum(b);
    r_b = blocks{b,3};
    
    idio_b = idio(:,ind_b); %restricting the analysis to the given block
    
    % block factor extraction with PC
    [v d] = eigs(cov(idio_b),r_b); %PCA
    fac_b = idio_b*v*diag(sqrt(1./diag(d))); %static projection
    
    tot_F(:,(r_cum_b+1):(r_cum_b+r_b)) = fac_b; % storing extracted fators in the relative columns
end

% defining the structure loadings selection matrix (here we impose the zero
% restrictions on factors that do not affect series)
lambda_struct = zeros(N,r_cum(end));
lambda_struct(:,1:r) = ones(size(lambda_struct(:,1:r))); % fill in th global factors with ones
for b = 1:size(blocks,1)
    ind_b = blocks{b,2};
    r_cum_b = r_cum(b);
    r_b = blocks{b,3};
    
    
    A = lambda_struct(ind_b,(r_cum_b+1):(r_cum_b+r_b));
    
    lambda_struct(ind_b,(r_cum_b+1):(r_cum_b+r_b)) = ones(size(A));
end

% start EB iterative algo
i = 0;
ll = [];
loglik = 0;
converged = 0;
while i < max_iter & ~converged
    i = i +1;
    
    lambda = ((tot_F'*tot_F)\tot_F'*Y)'; %OLS projection for loadings extaction from Lambda
    
    lambda = lambda.*lambda_struct; % set to zero following the structure
    
    
    tot_F = Y*pinv(lambda)'; % OLS projection for factor extractions
    tot_F = tot_F*diag(sqrt(1./diag(cov(tot_F)))); % factor rescaling
    
    S_lf = trace((Y - tot_F*lambda')*(Y - tot_F*lambda')'); % calculate RSS
    
    conv_lev = abs(S_lf - loglik)/abs(S_lf + loglik); % calculate conv. level
    
    converged = conv_lev < thresh; % store convergence status
    
    loglik = S_lf; % update loglik
    ll = [ll loglik]; % store squential loglik
end

facs = {tot_F(:,1:r)};
for b = 1:size(blocks)
    r_b = blocks{b,3};
    facs{b+1} = tot_F(:,r_cum(b)+1:r_cum(b)+r_b);
end

res.idio = Y-tot_F*lambda';
res.facs = facs;
res.factors = cell2mat(facs);
res.lambda = lambda;
res.F = tot_F;
res.loglik = ll;
res.iterations = i;

end
