%% FINAL VERSION - REPLICATION OF RESULTS IN THE PAPER - 21/05/2024
%% NB the recession indicator is state 1

clear all
clc
close all

%% load data
load Recessions_FRED
load('FRED_MD_04_2024.mat')

[T,N] = size(x);

x = ML_Standardize(x);

%% number of factors
r = 8; % as recommended by McCraken Ng 

%% EM parameters
%% best config up to 2024 heterosk 
homosked = 1; % or = 0
kmax = 100;
thresh = 10^(-4);
eps = [0 0.49];
min_eta = 10^(-50);

%% ESTIMATION
%% PCA with r factors
[V,D]=eigs(cov(x),r);
ghat = (D^(-1/2)*V'*x')';
Ahat = V*D^(1/2);
chihat = (Ahat*ghat')';
ehat = x-chihat;
clear V D

%% EM algorithm
%% pre-allocate
Phat = zeros(2,2,kmax+1);
rhohat = zeros(2*2,kmax+1);
Bhat = zeros(N,r,2,kmax+1);
Sigmaehat = zeros(N,N,2,kmax+1);
etahat = zeros(T,2,kmax+1);
llik = zeros(kmax+1,1);

xit1 = zeros(T,2,kmax+1);
xitt = zeros(T,2,kmax+1);
xitT = zeros(T,2,kmax+1);
xitT2 = zeros(T,2*2,kmax+1);


%% initialize parameters
Bhat(:,:,1,1) = Ahat;
Bhat(:,:,2,1) = Ahat;

Sigmaehat(:,:,1,1) = diag(diag(cov(ehat)));
Sigmaehat(:,:,2,1) = diag(diag(cov(ehat)));

Phat(:,:,1) = [0.5+eps(1,1) 1-0.5-eps(1,1);
    1-0.5-eps(1,2) 0.5+eps(1,2)];
xibar_init(1,:) = [(1-Phat(2,2,1)) (1-Phat(1,1,1))]./(2-Phat(1,1,1)-Phat(2,2,1))


rhohat(:,1) = ML_vec(Phat(:,:,1));
xi00 = [1 0];

%% initialize state
xit1(1,:,1) = (Phat(:,:,1)'*xi00')';

Ec = eye(2);

for t=1:T
    for j=1:2
        etahat(t,j,1) = (2*pi)^(-N/2)*(det(Sigmaehat(:,:,j,1)))^(-1/2)*...
            exp(-((x(t,:)-kron(Ec(j,:),ghat(t,:))*[Bhat(:,:,1,1) Bhat(:,:,2,1)]')*...
            Sigmaehat(:,:,j,1)^(-1)*...
            (x(t,:)-kron(Ec(j,:),ghat(t,:))*[Bhat(:,:,1,1) Bhat(:,:,2,1)]')')/2);
        if etahat(t,j,1)==0
            etahat(t,j,1)=min_eta;
        end
    end
end

for k=1:kmax
    fprintf('---EM iteration %d',k);
    llik(k,1) = sum(log(etahat(:,1,k)+etahat(:,2,k)));
    miscl(:,k) = log(etahat(:,1,k))-log(etahat(:,2,k));
    fprintf('--llik %f\n',llik(k,1));
    %% E-step
    %% Hamilton filter
    for t=1:T
        xitt(t,:,k) = (etahat(t,:,k).*xit1(t,:,k))/(ones(1,2)*(etahat(t,:,k).*xit1(t,:,k))');
        if t<T
            xit1(t+1,:,k) = (Phat(:,:,k)'*xitt(t,:,k)')';
        end
    end
    
    %% Kim smoother
    xitT(T,:,k) = xitt(T,:,k);
    xitT(T-1,:,k) = ((Phat(:,:,k)*(xitt(T,:,k)./xit1(T,:,k))').*xitt(T-1,:,k)')';
    
    revt = flipud((1:T-2)');
    
    for t=1:T-2
        xitT(revt(t),:,k) =  ((Phat(:,:,k)*(xitT(revt(t)+1,:,k)./xit1(revt(t)+1,:,k))').*xitt(revt(t),:,k)')';
    end
    
    for t=2:T
        xitT2(t,:,k) = kron(xitT(t-1,:,k),xitT(t,:,k));
    end
    
    if k>2 && abs(llik(k,1)-llik(k-1,1))/(0.5*abs(llik(k,1)+llik(k-1,1))) < thresh
        kstar(1,1) = k;
        break
    end
        
    %% M-step
    %% estimate parameters
    for j=1:2
        for t=1:T
            temp1(t,:,:) = x(t,:)'*ghat(t,:);
            temp2(t,:,:) = ghat(t,:)'*ghat(t,:);
        end
        
        Bhat(:,:,j,k+1) = squeeze(sum(temp1.*xitT(:,j,k),1))*(squeeze(sum(temp2.*xitT(:,j,k),1)))^(-1);
        
        clear temp1 temp2
        
        for t=1:T
            temp3(t,:) = (x(t,:)-ghat(t,:)*Bhat(:,:,j,k+1)').^2;
        end
          
        if homosked == 1
            Sigmaehat(:,:,j,k+1) = diag(diag(cov(ehat)));
        else
            Sigmaehat(:,:,j,k+1) = diag(sum(temp3.*xitT(:,j,k),1))./sum(xitT(:,j,k)); 
        end
        
        clear temp3
    end
    
    rhohat(:,k+1) = sum(xitT2(:,:,k),1)'./(kron(ones(2,1),sum(xitT(2:T,:,k),1)'));
    Phat(:,:,k+1) = reshape(rhohat(:,k+1),2,2);
    
    %% new initialization of state
    xit1(1,:,k+1) = mean(xitT(:,:,k),1);
    
    for t=1:T
        for j=1:2
            etahat(t,j,k+1) = (2*pi)^(-N/2)*(det(Sigmaehat(:,:,j,k+1)))^(-1/2)*...
                exp(-((x(t,:)-kron(Ec(j,:),ghat(t,:))*[Bhat(:,:,1,k+1) Bhat(:,:,2,k+1)]')*...
                Sigmaehat(:,:,j,k+1)^(-1)*...
                (x(t,:)-kron(Ec(j,:),ghat(t,:))*[Bhat(:,:,1,k+1) Bhat(:,:,2,k+1)]')')/2);
            if etahat(t,j,k+1)==0
                etahat(t,j,k+1)=min_eta;
            end
        end
    end
end

if k == kmax
    fprintf('EM did not converge\n');
    kstar = kmax;
end

xibar(1,:) = mean(xitT(:,:,kstar),1);
xibar(2,:) = [(1-Phat(2,2,kstar)) (1-Phat(1,1,kstar))]./(2-Phat(1,1,kstar)-Phat(2,2,kstar));
xibar(3,:) = mean(round(xitT(:,:,kstar)),1)


chihat0=[];
chihat0(:,:,1) = Bhat(:,:,1,kstar)*(ghat.*xitT(:,1,kstar))';
chihat0(:,:,2) = Bhat(:,:,2,kstar)*(ghat.*xitT(:,2,kstar))';


fhat(:,:,1) = xitT(:,1,kstar).*x*Bhat(:,1:r/2,1,kstar)/(Bhat(:,1:r/2,1,kstar)'*Bhat(:,1:r/2,1,kstar));
fhat(:,:,2) = xitT(:,2,kstar).*x*Bhat(:,r/2+1:r,2,kstar)/(Bhat(:,r/2+1:r,2,kstar)'*Bhat(:,r/2+1:r,2,kstar));

ftilde(:,:,1) = xitT(:,1,kstar).*x*Bhat(:,r/2+1:r,1,kstar)/(Bhat(:,r/2+1:r,1,kstar)'*Bhat(:,r/2+1:r,1,kstar));
ftilde(:,:,2) = xitT(:,2,kstar).*x*Bhat(:,1:r/2,2,kstar)/(Bhat(:,1:r/2,2,kstar)'*Bhat(:,1:r/2,2,kstar));


for j=1:r
    figure
    MB_FigureSize(25,15,1)
    hold all
    plot(DateN,(fhat(:,j)),'linewidth',2,'color','k')
    axis tight
    xticks(DateN([7:60:T]))
    xticklabels(datestr(DateN([7:60:T]),'yyyy'))
    box on
    grid on
end
% 
% for j=1:r
%     figure
%     MB_FigureSize(25,15,1)
%     hold all
%     plot(DateN,(ftilde(:,j)),'linewidth',2,'color','k')
%     axis tight
%     xticks(DateN([7:60:T]))
%     xticklabels(datestr(DateN([7:60:T]),'yyyy'))
%     box on
%     grid on
% end
% 

figure
MB_FigureSize(25,15,1)
hold all
plot(DateN,(xitT(:,1,kstar)),'linewidth',2,'color','k')
% plot(DateN,Recessions,'x','markersize',10,'color','b')
axis tight
xticks(DateN([7:60:T]))
xticklabels(datestr(DateN([7:60:T]),'yyyy'))
ylim([-0.1 1.1])
yticks([0 1])
yticklabels([{'0'} {'1'}])
box on
grid on

figure
MB_FigureSize(25,15,1)
hold all
plot(DateN,(xitT(:,2,kstar)),'linewidth',2,'color','k')
% plot(DateN,Recessions,'x','markersize',10,'color','b')
axis tight
xticks(DateN([7:60:T]))
xticklabels(datestr(DateN([7:60:T]),'yyyy'))
ylim([-0.1 1.1])
yticks([0 1])
yticklabels([{'0'} {'1'}])
box on
grid on

% figure
% MB_FigureSize(25,15,1)
% hold all
% plot(llik(1:kstar,1),'linewidth',2,'color','k')
% xlabel('$k$','interpreter','latex','fontsize',15)
% axis tight
% box on
% grid on

%% degree of concordance Harding and Pagan

Ind = round(xitT(:,1,kstar));
type1=0; type2=0;
good1=0; good2=0;

for t=1:T
    if Ind(t,1)== 1 & Recessions(t,1)==0
        type2 = type2+1;
    elseif Ind(t,1)== 0 & Recessions(t,1)==1
        type1 = type1+1;
    elseif Ind(t,1)== 1 & Recessions(t,1)==1
        good1 = good1+1;
    elseif Ind(t,1)== 0 & Recessions(t,1)==0
        good2 = good2+1;
    end
end

doc = (good1+good2)/T;

FP = type2/T;

FN = type1/T;



