function y = EKF(pars, Y, rlb, mats, aJ, bJ, sigmasJ, index)
% Extended Kalman Filter to update beta_t|t and to generate log likelihood
% INPUT:
%       Y, Txj yields
%       mats, maturities in terms of year
%       dt:  length of each period in basis of year. 1/100 year
%       rlb:  lower bound rate, typically 0 
% possible OUTPUT:
%         llf: scalar, -loglikelihood
%         SR: tx1, annualized short rate 
%         Xf: 4xt filtered variables [\tau \tilde(X)] 
%         XT: 4xt generated state variables

T = size(Y,1) ;
lenJ = length(mats);
N = 3 ;
Y = Y';
% parameterization
Phi = pars.Phi;
omega = pars.omega;
sigtau2 = pars.sigtau2;
a = pars.a;
p = pars.p;
gamma = [1-sum(a); a];
Xbar = [-sum(p); p];
sige2 = pars.sige2;
R = sige2 * eye(lenJ) ;
mu_Z = [0; (eye(N)-Phi)*Xbar] ;
Phi_Z = [1 zeros(1,N); (eye(N)-Phi)*gamma Phi] ;
Q = [sigtau2 sigtau2*gamma'; sigtau2*gamma omega] ;
% Initialization
X1 = NaN(N+1,T+1); % E_t-1[X_t]; t = 1,...,T+1
X2 = NaN(N+1,T+1); % E_t-1[X_t-1]; t = 1,...,T+1; only meanningful after t>=2
V1 = NaN((N+1)*(T+1),N+1); % COV_t-1[X_t]; t = 1,...,T+1
V2 = NaN((N+1)*(T+1),N+1); % COV_t-1[X_t-1]; t = 1,...,T+1; only meanningful after t>=2
loglikvec = NaN(1,T); % lf(t); t = 1,...,T

X1(:,1) = zeros(4,1); 
V1(1:N+1,:) = 100*eye(N+1) ; 
% X1(1,1) = 6;

% Filtering the shadow rate
for t = 1:1:T
    musJ = aJ + bJ*X1(2:N+1,t);  % (K-1)x1
    z1_temp = (musJ-rlb)./sigmasJ;
    z2_temp = rlb + (musJ-rlb).*normcdf(z1_temp) + sigmasJ.*normpdf(z1_temp);
    h = normcdf(z1_temp) ;  
    h = [h h h].*bJ ;         % Kx3 
    H = [zeros(lenJ,1) h];    % Kx4
    err = Y(:,t) - z2_temp;    % jx1
    S = H*V1((t-1)*4+1:t*4,:)*H' + R;   % jxj
    invS = S\eye(size(S,1)) ;
    K = V1((t-1)*4+1:t*4,:)*H'*invS ;   % 4xj
    % Calculate the likelihood
    loglikvec(t) = -1/2 * (lenJ*log(2*pi) + log(det(S)) + err'*invS*err) ;
    % Update
    X2(:,t+1) = X1(:,t) + K*err;
    V2(4*t+1:4*(t+1),:) = (eye(N+1)-K*H)*V1((t-1)*4+1:t*4,:);
    % Predict
    X1(:,t+1) = mu_Z + Phi_Z * X2(:,t+1);
    V1(4*t+1:4*(t+1),:) = Phi_Z * V2(4*t+1:4*(t+1),:) * Phi_Z' + Q;
end
llf = sum(loglikvec(3:end)); 
SR = sum(X2(2:end,2:end),1) ;   % 1xT
Xf = X2(:,2:end);
if index==0
    y = -llf ;
elseif index==1
    y = SR;
elseif index==2
    y = Xf ;
else
    % generating sequence of state variables (beta1 beta2,...,betaT)
    % based on Kim and Nelson eq (8.9), (8.16), (8.17)
    stt = X2(:,2:end) ;
    pttb = V2(N+2:end,:) ;
%     muT = X2(:,end) ;
%     PT = V2(:,:,end) ;
%     XT = nan(N+1,T) ;
%     XT(:,T) =  mvnrnd(muT,PT,1) ;
    
    sTT = X2(:,end);
    pTT = V2(end-N:end,:);
    
    [pchol, Sflag2] = chol(pTT);
    if Sflag2 ~= 0
        [UP, SP, ~] = svd(pTT);
        pchol = (UP*sqrt(SP))';
    end
    sT = sTT + pchol'*randn(4,1);
    s = zeros(N+1,T);
    s(:,T) = sT;
    
    for jj = 1:T-1
        sTm1 = stt(:,T-jj); %sTm1=st|t
%         pTm1 = pttb(:,:,T-jj); %pTm1=pt|t
        pTm1 = pttb(end-4*jj-N:end-jj*4,:); %pTm1=pt|t
        J = pTm1*Phi_Z'/(Phi_Z*pTm1*Phi_Z'+Q);
        sTm1T = sTm1+J*(sT-mu_Z-Phi_Z*sTm1);   %sTm1T=st|t,st+1
        pTm1T = pTm1-J*Phi_Z*pTm1;             %pTm1T=pt|t,st+1
        [pchol, Sflag2] = chol(pTm1T);
        if Sflag2 ~= 0 
           [UP, SP, ~] = svd(pTm1T);
           pchol = (UP*sqrt(SP))';
        end
        sT = sTm1T+pchol'*randn(4,1);
        s(:,T-jj) = sT;
    end
    y = s;
end



end


        