%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% A function perfoming the estimation a block DFM
% Imputs:
% 1) Matrix of data -> X;
% 2) P -> struct with the global parameters;
% 3) blocks -> cell-array with the block structure info

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


function Res = EM_BDFM(X,P,blocks,blocks_full,conv_ind)

if nargin < 5
conv_ind = 1:10;    
end    

if nargin < 4
blocks_full = blocks;
end

% X: T x n panel of data
% P: structure containing settings of the model and estimation

thresh = P.conv_thresh;
r = P.r; % number of factors
p = P.p; % number of lags in the factor VAR
max_iter = P.max_iter; % maximal number of EM iterations

%--------------------------------------------------------------------------
% Preparation of the data
%--------------------------------------------------------------------------
[T,N] = size(X)

% Standardise the time series
Mx = nanmean(X); % read means
Wx = nanstd(X); % read variances
xNaN = (X-repmat(Mx,T,1))./repmat(Wx,T,1); %standardize

%--------------------------------------------------------------------------
% Initialization through Eickmeier and Breitung algorithm
%--------------------------------------------------------------------------
% Initialization is performed with a completed matrix. This matrix is
% filled within the initilization function

optNaN.method = 1; % Set the method for NaNs removal
optNaN.k = 5; %order of the moving average to smooth the series after removing the missing observations

[A, C, Q, R, Z_0, V_0] = EB_init(xNaN,r,p,optNaN,blocks,thresh);


%--------------------------------------------------------------------------
% Model estimation
%--------------------------------------------------------------------------

% some auxiliary variables for the iterations
previous_loglik = -inf;
num_iter = 0; % initialize iterations counter
LL = -inf; % initialize loglik vector
converged = 0; % initialize convergence condition (0 = not convderged)

% Store the transpose  matrix y for the estimation WITH missing data
y = xNaN';


%--------------------------------------------------------------------------
%THE EM LOOP
%--------------------------------------------------------------------------

%The model can be written as
% y = C*Z + e;
% Z = A*Z(-1) + v
% R = cov(e)
% Q = cov(v)
%where y is NxT, Z is (pr)xT, etc


r_tot = r + sum(cell2mat(blocks(:,3))); % read the total number of factors (glocl + local factors)


while (num_iter < max_iter) & ~converged % condition to keep iterating
    
    % Print the iteration data to keep track of the program advancements
    fprintf('\n')
    disp(strcat("Iteration number ", string(num_iter)))
    fprintf('\n')
    
    % At each cycle perform the E-step and the M-step
    [C_new, R_new, A_new, Q_new, Z_0, V_0, loglik] = EMstep_block(y, A, C, Q, R, Z_0, V_0, r, r_tot, blocks);
    
    % Update the model parameters
    C(:,1:r_tot) = C_new;
    R = R_new;
    A = A_new;
    Q = Q_new;
    
    % Checking convergence
    [converged,decrease(num_iter+1)] = em_converged2(loglik, previous_loglik, thresh,1);

    LL = [LL loglik]; % store loglikelihood progression into the array LL
    previous_loglik = loglik; % update the loglik
    num_iter =  num_iter + 1; % increase the iteration number
   
end



% Final run of the Kalman filter for the factors
[Zsmooth Vsmooth] = runKF_hp(y, A, C, Q, R, Z_0, V_0);
Zprint = Zsmooth';
F = Zprint(2:end,1:r_tot);
Res.x_sm = Zprint(2:end,:)*C';
Res.X_sm = repmat(Wx,T,1).*Res.x_sm+repmat(Mx,T,1);
Res.F = F;


% calculating variances of loadings

[C_var] = var_calc(y, C, R, Zsmooth, Vsmooth, r, blocks);

% final regression extracting the loadings for all the flows given the
% dataset

%[C C_var R] = complete_reg(y, A_new, Q_new, C, R, Zsmooth, Vsmooth, r, blocks, blocks_full, conv_ind);



%--------------------------------------------------------------------------
%   Loading the structure with the results
%--------------------------------------------------------------------------
Res.C = C;
Res.C_var = C_var;
Res.R = R;
Res.A = A;
Res.Q = Q;
Res.Z_0 = Z_0;
Res.V_0 = V_0;
Res.r = r;
Res.p = p;
Res.Mx = Mx;
Res.Wx = Wx;
Res.LL = LL;


Res

% Res.loglik   = LL;
% Res.num_iter = num_iter;
% Res.X_sm = X_sm;
% Res.converge = converged;
% Res.decrease = any(decrease);
% Res.decr_iter = min(find(decrease));
% decrease = any(decrease);
%--------------------------------------------------------------------------
%PROCEDURES
%--------------------------------------------------------------------------

end

function  [C_new, R_new, A_new, Q_new, Z_0, V_0, loglik] = EMstep_block(y, A, C, Q, R, Z_0, V_0, r, r_tot ,blocks)

[n,T] = size(y);

% Compute the (expected) sufficient statistics for a single Kalman filter sequence.

%Running the Kalman filter with the current estimates of the parameters
[Zsmooth, Vsmooth, VVsmooth, loglik] = runKF_hp(y, A, C, Q, R, Z_0, V_0);


EZZ = Zsmooth(:,2:end)*Zsmooth(:,2:end)'+sum(Vsmooth(:,:,2:end),3);          %E(Z'Z)
EZZ_BB = Zsmooth(:,1:end-1)*Zsmooth(:,1:end-1)'+sum(Vsmooth(:,:,1:end-1),3); %E(Z(-1)'Z_(-1))
EZZ_FB = Zsmooth(:,2:end)*Zsmooth(:,1:end-1)'+sum(VVsmooth,3);               %E(Z'Z_(-1))


A_new = A;
A_new(1:r,1:r) = EZZ_FB(1:r,1:r) * inv(EZZ_BB(1:r,1:r)); % common component of A

Q_new = Q;
Q_new(1:r,1:r) = (EZZ(1:r,1:r) - A_new(1:r,1:r)*EZZ_FB(1:r,1:r)') / T;

C_new = zeros(size(C));
r_cum = r;
n_cum = 0;


%disp('Loading matrix update..')

for i=1:size(blocks,1)
    
    r_b = blocks{i,3};
    blocks{i,5} = r_cum;
    r_cum = r_cum + r_b;
end
ns = size(blocks,1);
for ss=1:ns

    
    r_b = blocks{ss,3};
    ind_b = blocks{ss,2};
    n_b = length(ind_b);
    r_cum = blocks{ss,5};
    
    
    ind_block_b = (r_cum+1):(r_cum + r_b);
    
    A_new(ind_block_b,ind_block_b) = EZZ_FB(ind_block_b,ind_block_b) * inv(EZZ_BB(ind_block_b,ind_block_b));
    Q_new(ind_block_b,ind_block_b) = (EZZ(ind_block_b,ind_block_b) - A_new(ind_block_b,ind_block_b)*EZZ_FB(ind_block_b,ind_block_b)') / T;
    
    
    %E(Y'Y) & E(Y'Z)
    y_b = y(ind_b,:);
    nanYb = isnan(y_b);
    y_b(nanYb) = 0;
    
    
    %%%%% UPDATE OF THE COEFFICIENTS MATRIX C BLOCK BY BLOCK
    

    denom_size = n_b*(r+r_b);
    denom = spalloc(denom_size, denom_size, n_b*(r+r_b));
    
    nom = zeros(n_b,r+r_b);
    
    r_cb_ind = [1:r,r_cum+1:r_cum+r_b];
    
    nanYb = sparse(nanYb);
    
    for t=1:T        
        nanYt = diag(~nanYb(:,t));
    
  
        aux = sparse(Zsmooth(r_cb_ind,t+1)*Zsmooth(r_cb_ind,t+1)'+Vsmooth(r_cb_ind,r_cb_ind,t+1));                       
        aux2 = matlab.internal.sparse.kronSparse(aux, nanYt);        
        denom = denom + aux2; 
        
        aux4 = y_b(:,t)*Zsmooth(r_cb_ind,t+1)';
        nom = nom + aux4;
    end

    vec_C = denom\nom(:);
    C_par = reshape(vec_C,n_b,(r+r_b));
    C_new(ind_b,1:r) = C_par(:,1:r);
    C_new(ind_b,r_cum+1:r_cum+r_b) = C_par(:,r+1:end);
    
    
end

%fprintf('\n');

nanY = isnan(y);
y(nanY) = 0;


clear vec_C
clear denom
clear nom
clear C


%disp('Idiosyncratic covariance matrix update..')

R_new = sparse(zeros(n,1));
for t=1:T
    %   fprintf(1,'\b%d',t/T);
    
    yC = (~nanY(:,t)).*C_new;
    
    R_new = R_new + (y(:,t)-yC*Zsmooth(1:r_tot,t+1)).^2 ...
        + sum((yC*Vsmooth(1:r_tot,1:r_tot,t+1)).*yC,2)...
        + R'.*(nanY(:,t));
    
end

fprintf('\n');

R_new = R_new'/T;


% Initial conditions
Z_0 = Zsmooth(:,1); % randn(size(Zsmooth,1),1); %
V_0 = Vsmooth(:,:,1);

end

%--------------------------------------------------------------------------

function [converged, decrease] = em_converged2(loglik, previous_loglik, threshold, check_increased)
% EM_CONVERGED Has EM converged?
% [converged, decrease] = em_converged(loglik, previous_loglik, threshold)
%
% We have converged if the slope of the log-likelihood function falls below 'threshold',
% i.e., |f(t) - f(t-1)| / avg < threshold,
% where avg = (|f(t)| + |f(t-1)|)/2 and f(t) is log lik at iteration t.
% 'threshold' defaults to 1e-4.
%
% This stopping criterion is from Numerical Recipes in C p423
%
% If we are doing MAP estimation (using priors), the likelihood can decrase,
% even though the mode of the posterior is increasing.

if nargin < 3, threshold = 1e-4; end
if nargin < 4, check_increased = 1; end

converged = 0;
decrease = 0;

if check_increased
    if loglik - previous_loglik < -1e-3 % allow for a little imprecision
        fprintf(1, '******likelihood decreased from %6.4f to %6.4f!\n', previous_loglik, loglik);
        decrease = 1;
    end
end

delta_loglik = abs(loglik - previous_loglik);
avg_loglik = (abs(loglik) + abs(previous_loglik) + eps)/2;
if (delta_loglik / avg_loglik) < threshold, converged = 1; end

end


function [ A, C, Q, R, chi, res, initZ, initV] = PCA_proj(x,r,p,optNaN)

OPTS.disp=0;

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

[T,N] = size(x);
% Eigenval decomp of cov(x) = VDV', only r largest evals
[ v, d ] = eigs(cov(xBal),r,'lm',OPTS);
v = v*diag(sqrt(N)*ones(r,1));
d_inv = diag(1./diag(d));

% Static predictions
chi = xBal*v*d_inv*v';

res = xBal-chi;
res(indNaN) = nan;

% Observation equation
R = diag(nanvar(res));

% Observation equation
C = [v zeros(N,r*(p-1))];

% Transition equation
% Estimate A & Q from stacked F(t) = A*F(t-1) + e(t);
F = xBal*v*d_inv;



z = F;
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 = zeros(r*p,r*p)';
A_temp = inv(Z'*Z)*Z'*z;
A(1:r,1:r*p) = A_temp';
A(r+1:end,1:r*(p-1)) = eye(r*(p-1));

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

rp2 = (r*p)^2;
% Initial conditions
initZ = zeros(size(Z,2),1); % %[randn(1,r*(nlag+1))]';
initV = reshape(inv(eye(rp2)-kron(A,A))*Q(:),r*p,r*p);

end

%% EB initialization functions
function [A, C, Q, R, initZ, initV] = EB_init(x,r,p,optNaN,blocks,thresh)


OPTS.disp=0;

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

[T,N] = size(x);
max_iter = 100;


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))]';

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 projection 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.lambda = lambda;
res.F = tot_F;
res.loglik = ll;

end

%% Blockwise regression for loadings calculation

function [C_new C_var R_new] = complete_reg(y, A_new, Q_new, C, R, Zsmooth, Vsmooth, r, blocks, blocks_full, conv_ind)

%disp('Loading matrix update..')

v = matfile('full_data.mat'); % uploading the matfile containing the full dataset

[N T] = size(v.y_f)


C_new = zeros(N,size(C,2));
r_cum = r;
n_cum = 0;


for i=1:size(blocks_full,1)
    
    r_b = blocks_full{i,3};
    blocks_full{i,5} = r_cum;
    r_cum = r_cum + r_b;
end
r_tot = r_cum;

% block by block calculations of coefficients

ns = size(blocks_full,1);
for ss=1:ns
        
    r_b = blocks_full{ss,3};
    ind_b = blocks_full{ss,2};
    
    n_b = length(ind_b);
    r_cum = blocks_full{ss,5};
    
    ind_block_b = (r_cum+1):(r_cum + r_b);
    
    %E(Y'Y) & E(Y'Z)
    y_b = v.y_f(ind_b,:);
    nanYb = isnan(y_b);
    y_b(nanYb) = 0;
    
    
    %%%%% UPDATE OF THE COEFFICIENTS MATRIX C BLOCK BY BLOCK
    
    %     % release 1.0
    %     denom = zeros(n_b*(r+r_b),n_b*(r+r_b));
    
    %     % release 2.0
    %     denom = sparse(zeros(n_b*(r+r_b),n_b*(r+r_b)));
    
    % release 3.0
    denom_size = n_b*(r+r_b);
    denom = spalloc(denom_size, denom_size, denom_size);
    
    nom = zeros(n_b,r+r_b);
    
    r_cb_ind = [1:r,r_cum+1:r_cum+r_b];
    
    
    nanYb = sparse(nanYb); % do it once, at the beginning! KUKOC_TRICK
    
    for t=1:T
        %nanYt = sparse(diag(~nanYb(:,t)));
        nanYt = diag(~nanYb(:,t));
        
        
        aux = sparse(Zsmooth(r_cb_ind,t+1)*Zsmooth(r_cb_ind,t+1)'+Vsmooth(r_cb_ind,r_cb_ind,t+1));
        
        % size(aux) = 2 x 2
        % size(nanYt) = 159 x 159
        
        aux2 = matlab.internal.sparse.kronSparse(aux, nanYt);
        %         aux3 = kronmult(aux,nanYt);
        denom = denom + aux2;
        
        aux4 = y_b(:,t)*Zsmooth(r_cb_ind,t+1)';
        nom = nom + aux4;
        
        % size(aux2) = 159 x 2
        % size(nome) = 159 x 2
        
    end
    
    
    if (r+r_b == 2)
        Aa = diag(denom(1:n_b,1:n_b));
        Bb = diag(denom(1:n_b,n_b+1:(r+r_b)*n_b));
        Cc = diag(denom(n_b+1:(r+r_b)*n_b,1:n_b));
        Dd = diag(denom(n_b+1:(r+r_b)*n_b,n_b+1:(r+r_b)*n_b));
        
        vec_C =  block_str_inv(Aa,Bb,Cc,Dd)*nom(:);
    else
        vec_C = denom\nom(:);
    end
    
    C_par = reshape(vec_C,n_b,(r+r_b));
    C_new(ind_b,1:r) = C_par(:,1:r);
    C_new(ind_b,r_cum+1:r_cum+r_b) = C_par(:,r+1:end);
    
end

C_new = C_new(conv_ind,:);
%fprintf('\n');

[n,T] = size(y);

nanY = isnan(y);
y(nanY) = 0;


R_new = sparse(zeros(n,1));
for t=1:T
    
    yC = (~nanY(:,t)).*C_new;
    
    R_new = R_new + (y(:,t)-yC*Zsmooth(1:r_tot,t+1)).^2 ...
        + sum((yC*Vsmooth(1:r_tot,1:r_tot,t+1)).*yC,2)...
        + R'.*(nanY(:,t));
    
end

fprintf('\n');

R_new = R_new'/T;

r_cum = r;
n_cum = 0;

for i=1:size(blocks,1)
    
    r_b = blocks{i,3};
    blocks{i,5} = r_cum;
    r_cum = r_cum + r_b;
end
r_tot = r_cum;

% block by block calculations of coefficients variances
C_var = zeros(size(C_new));
R_new = sparse(R_new);

ns = size(blocks,1);
for ss=1:ns
    
    %fprintf(1,'\b%d',100*round(ss/size(blocks,1),6));
    
    r_b = blocks{ss,3};
    ind_b = blocks{ss,2};
    n_b = length(ind_b);
    r_cum = blocks{ss,5};
    
    ind_block_b = (r_cum+1):(r_cum + r_b);
    
    %E(Y'Y) & E(Y'Z)
    y_b = y(ind_b,:);
    nanYb = isnan(y_b);
    y_b(nanYb) = 0;
    
    %%%%% UPDATE OF THE COEFFICIENTS MATRIX C BLOCK BY BLOCK
    
    denom = sparse(n_b*(r+r_b),n_b*(r+r_b));
    
    r_cb_ind = [1:r,r_cum+1:r_cum+r_b];
    nanYb = sparse(nanYb);
    for t=1:T
        iRnanYt = 2*(1./R_new(:,ind_b).*diag(~nanYb(:,t)));
        
        aux = sparse(Zsmooth(r_cb_ind,t+1)*Zsmooth(r_cb_ind,t+1)'+Vsmooth(r_cb_ind,r_cb_ind,t+1));
        aux2 = matlab.internal.sparse.kronSparse(aux,iRnanYt);
        
        denom = denom + aux2;
    end
    
    
    if (r+r_b == 2)
        Aa = diag(denom(1:n_b,1:n_b));
        Bb = diag(denom(1:n_b,n_b+1:(r+r_b)*n_b));
        Cc = diag(denom(n_b+1:(r+r_b)*n_b,1:n_b));
        Dd = diag(denom(n_b+1:(r+r_b)*n_b,n_b+1:(r+r_b)*n_b));
        
        iA = 1./Aa;
        center = 1./(Dd - Cc.*iA.*Bb);
        
        A1 = iA + iA.*Bb.*center.*Cc.*iA;
        C1 = -center.*Cc.*iA;
        
        C_par_var = [A1 C1];
    else
        
        vec_C = inv(denom);
        C_par_var = reshape(vec_C,n_b,(r+r_b));
        
    end
    C_var(ind_b,1:r) = C_par_var(:,1:r);
    C_var(ind_b,r_cum+1:r_cum+r_b) = C_par_var(:,r+1:end);
end

end



function [C_var] = var_calc(y, C, R, Zsmooth, Vsmooth, r, blocks)

T = size(Zsmooth,2)-1

% block by block calculations of coefficients variances
C_var = sparse(size(C));
R_new = sparse(R);

r_cum = r;
n_cum = 0;

for i=1:size(blocks,1)
    
    r_b = blocks{i,3};
    blocks{i,5} = r_cum;
    r_cum = r_cum + r_b;
end
r_tot = r_cum;


ns = size(blocks,1);
for ss=1:ns
    
    %fprintf(1,'\b%d',100*round(ss/size(blocks,1),6));
    
    r_b = blocks{ss,3};
    ind_b = blocks{ss,2};
    n_b = length(ind_b);
    r_cum = blocks{ss,5};
    
    ind_block_b = (r_cum+1):(r_cum + r_b);
    
    %E(Y'Y) & E(Y'Z)
    y_b = y(ind_b,:);
    nanYb = isnan(y_b);
    y_b(nanYb) = 0;
    
    %%%%% UPDATE OF THE COEFFICIENTS MATRIX C BLOCK BY BLOCK
    
    denom = sparse(n_b*(r+r_b),n_b*(r+r_b));
    
    r_cb_ind = [1:r,r_cum+1:r_cum+r_b];
    nanYb = sparse(nanYb);
    for t=1:T
        iRnanYt = 2*(1./R_new(:,ind_b).*diag(~nanYb(:,t)));
        
        aux = sparse(Zsmooth(r_cb_ind,t+1)*Zsmooth(r_cb_ind,t+1)'+Vsmooth(r_cb_ind,r_cb_ind,t+1));
        aux2 = matlab.internal.sparse.kronSparse(aux,iRnanYt);
        
        denom = denom + aux2;
    end
    
    
    if (r+r_b == 2)
        Aa = diag(denom(1:n_b,1:n_b));
        Bb = diag(denom(1:n_b,n_b+1:(r+r_b)*n_b));
        Cc = diag(denom(n_b+1:(r+r_b)*n_b,1:n_b));
        Dd = diag(denom(n_b+1:(r+r_b)*n_b,n_b+1:(r+r_b)*n_b));
        
        iA = 1./Aa;
        center = 1./(Dd - Cc.*iA.*Bb);
        
        A1 = iA + iA.*Bb.*center.*Cc.*iA;
        C1 = -center.*Cc.*iA;
        
        C_par_var = [A1 C1];
    else
        
        vec_C = inv(denom);
        C_par_var = reshape(vec_C,n_b,(r+r_b));
        
    end
    C_var(ind_b,1:r) = C_par_var(:,1:r);
    C_var(ind_b,r_cum+1:r_cum+r_b) = C_par_var(:,r+1:end);
end

end



function [L] = block_str_inv(A,B,C,D)

%This function takes 4 vectors, namely the diagonals of the four
%blocks and compute the inverse
n = length(A);

iA = 1./A;
center = 1./(D - C.*iA.*B);

A1 = iA + iA.*B.*center.*C.*iA;
B1 = -iA.*B.*center;
C1 = -center.*C.*iA;
D1 = center;

L = [sparse(1:n,1:n,A1) sparse(1:n,1:n,B1)
    sparse(1:n,1:n,C1) sparse(1:n,1:n,D1)];


end











