% This file solve for B in a model with lagged expectations
% Written by Yossi Yakhin: yossiya@bgu.ac.il

function B = solution(B_guess);

global N m n A A0 A1tilde Btilde BETA OMEGA THETA PSI

for i=1:N
    BB(:,:,i)             = B_guess(:,(i-1)*m+1:i*m);
    Bcol((i-1)*n+1:i*n,:) = BB(:,:,i);
    GG(:,:,i)             = OMEGA(:,(i-1)*n+1:i*n);
end

% Apower is a stricktly lower block-triangular matrix, each block is nxn.
% Each block along the diagonal -i (i=1 to N-1) contains the matrix A^i.
Apower  = zeros(N*n,N*n);
for i=1:N-1
    Apower = Apower + kron(diag(ones(N-i,1),-i),A^i);
    i
end

% Apower1 is similar to Apower but with A^0 along the main diagonal 
Apower1 = eye(N*n) + Apower;


iter = 0;
crit = 1;
while crit>1e-12
    
    % Updating B_0 through B_N-1
    BB(:,:,1)   = inv(A0*A + A1tilde - GG(:,:,1))*(-A0*BB(:,:,2) - BETA);
    Bcol(1:n,:) = BB(:,:,1);
    for j=2:N-1
        BB(:,:,j) = inv(A0*A + A1tilde - GG(:,:,j))*...
                    (-A0*BB(:,:,j+1) + GG(:,:,j)*Apower((j-1)*n+1:j*n,:)*Bcol);
        Bcol((j-1)*n+1:j*n,:) = BB(:,:,j);
    end
    BB(:,:,N) = inv(A0*A + A1tilde - GG(:,:,N))*GG(:,:,N)*Apower((N-1)*n+1:N*n,:)*Bcol;
    Bcol((N-1)*n+1:N*n,:) = BB(:,:,N);

    % Constructing the B matrix
    for i=1:N
        B(:,(i-1)*m+1:i*m) = BB(:,:,i);
    end

    % Constructing the PSI matrix
    PHIcol = Apower1*Bcol;
    PSI = zeros(N*n,N*m);
    for i=1:N
        PSI((i-1)*n+1:i*n,(i-1)*m+1:i*m) = PHIcol((i-1)*n+1:i*n,:);
    end
    
    % Evaluating equation (6) in the notes
    errors = (A0*A + A1tilde)*B + A0*B*THETA + Btilde - OMEGA*PSI;
    crit = max(max(abs(errors)))
    iter = iter + 1
end