%{
    Xavier Martin G. Bautista
    Steven G. Craig
    Annie Yu-Hsin Hsu
    Bent E. Sorensen
    Vasundhara Tanwar
    Priyam Verma 

    kalman.m
    -------------

    This code contains everything related to the Kalman filter and 
    smoother.    

%}
classdef kalman
    methods(Static)
        %% 1. Clean the data.
        
        function dat = cleaner(datain)
            
            dat = struct();

            dat.panelid = xlsread(datain,1,'A:A'); 
            dat.timeid = xlsread(datain,1,'B:B'); 
            
            dat.N = max(dat.panelid);                                       % States.
            dat.T = max(dat.timeid);                                        % Years

            y = log(xlsread(datain,1,'C:C'));                               % Income per capita.
            
            dat.y = reshape(y,dat.T,dat.N);
            dat.dy = dat.y(2:end,:)-dat.y(1:end-1,:);
                     
            dat.T = size(dat.dy,1);                                         % Time periods in differences.
            dat.g = mean(dat.dy);
            dat.dy = dat.dy-dat.g;                                          % Remove average over time for each state.
            
            dat.y0 = dat.y(1,:);
            
        end
        
        %% 2. Kalman Filter.
        
        function LL = filter(sigma0,dat)
            %% Set up model and data.
            
                dy = dat.dy';
            
                m = 3;                                                      % Number of state variables.
                
                K = zeros(m,m);                                             % Matrix of constraints.
                K(2,3) = 1;
                
                R = zeros(3,2);
                R(1,1) = 1;
                R(2,2) = 1;
                R(3,2) = -1;


                Z = [1; 1; 0];                                              % For measurement equation.

                Q = zeros(2,2);                                             % Variance matrix of shocks.
                Q(1,1) = sigma0(1)^2;
                Q(2,2) = sigma0(2)^2;
                
            %% Loglikelihood function.    
                
                N = dat.N;
                T = dat.T;
                
                LL = 0;
                               
                for i = 1:N                                                 % Loop over states.
                    
                    X = zeros(m,1);                                         % alpha(t|t)

                    if dat.pinit == 1
                        RQR = reshape(R*Q*R',[],1);
                        KK = inv(eye(m*m)-kron(K,K));
                        P = reshape(KK*RQR,m,m);
                    else
                        P = eye(m,m);
                    end

                    for j = 1:T
                        % Prediction equations. 
                        
                        Xhat = K*X;                                         % alpha(t|t-1)
                        ypred = Z'*Xhat;                                    % Y(t|t-1)
                        Phat = K*P*K' + R*Q*R';                             % P(t|t-1)

                        % Error and variance matrix.
                        
                        nu = dy(i,j) - ypred;     
                        F = Z'*Phat*Z;
                        
                        % Loglikelihood.
                        
                        LL = LL - 0.5*log(F) - 0.5*nu'*inv(F)*nu... 
                                  - 0.5*log(2*pi);
                              
                        % Updating equations.
                        
                        X = Xhat + Phat*Z*inv(F)*nu;                        % alpha(t|t)
                        P = Phat - Phat*Z*inv(F)*Z'*Phat';                  % P(t|t)
                    end

                end
                                
                LL = -LL;
            
        end
        
        %% 3. Kalman Smoother.
        
        function dat = smoother(sigma0,dat)
            %% Set up model.
            
                dy = dat.dy';
            
                m = 3;                                                      % Number of state variables.
                K = zeros(m,m);
                K(2,3) = 1;

                Z = [1; 1; 0];

                R = zeros(m,2);
                R(1,1) = 1;
                R(2,2) = 1;
                R(3,2) = -1;

                Q = zeros(2,2);                                             % Variance matrix of shocks.
                Q(1,1) = sigma0(1)^2;
                Q(2,2) = sigma0(2)^2;
                
                Pmat     = zeros(m,dat.T*m);    
                Pmat_hat = zeros(m,dat.T*m);   

                Xmat      = zeros(m,dat.T);                                 % In the notes, these are alpha.
                Xmat_hat  = zeros(m,dat.T);
                Xsmat     = zeros(m,dat.T);                                 
                Xsmat_hat = zeros(m*dat.N,dat.T);
                
            %% Begin Kalman filter.   
            
                N = dat.N;
                T = dat.T;
                
                for i = 1:N                                                 % Loop over states.
                    
                    X = zeros(m,1);

                    if dat.pinit == 1
                        RQR = reshape(R*Q*R',[],1);
                        KK = inv(eye(m*m)-kron(K,K));
                        P = reshape(KK*RQR,m,m);
                    else
                        P = eye(m,m);
                    end

                    for j = 1:T
                        % Prediction equations. 
                        
                        Xhat = K*X;                                         % alpha(t|t-1)
                        ypred = Z'*Xhat;                                    % Y(t|t-1)
                        Phat = K*P*K' + R*Q*R';                             % P(t|t-1)

                        % Error and variance matrix.
                        
                        nu = dy(i,j) - ypred;     
                        F = Z'*Phat*Z;
                              
                        % Updating equations.
                        
                        X = Xhat + Phat*Z*inv(F)*nu;                        % alpha(t|t)
                        P = Phat - Phat*Z*inv(F)*Z'*Phat';                  % P(t|t)
                        
                        Xmat(:,j) = X;                                      % Save alpha(t|t).
                        Xmat_hat(:,j) = Xhat;                               % Save alpha(t|t-1).
                          
                        Pmat(:,(j-1)*3+1:(j-1)*3+3) = P;                    % Save P(t|t).
                        Pmat_hat(:,(j-1)*3+1:(j-1)*3+3) = Phat;             % Save P(t|t-1).
              
                    end
                    
                %% Begin smoother.
                    
                    Xsmat(:,T) = Xmat(:,T);                                 % Last period.
                    
                    P    = Pmat(:,3*T-5:3*T-3);                             % P(t-1|t-1)
                    Phat = Pmat_hat(:,3*T-2:3*T);                           % P(t|t-1)
                    Jt   = P*K'*inv(Phat);             

                    Xsmat(:,T-1) = Xmat(:,T-1)...                           % alpha(t|T)
                                           + Jt*(Xsmat(:,T)-Xmat_hat(:,T));

                    for n = 2:T-1                                           % Recursion.
                        P    = Pmat(:,3*T-n*3-2:3*T-n*3);
                        Phat = Pmat_hat(:,3*T-n*3+1:3*T-n*3+3);  
                        Jt   = P*K'*inv(Phat);  
                        
                        Xhat = Xmat_hat(:,T-n+1);
                        Xsmat(:,dat.T-n) = Xmat(:,T-n)...
                                                + Jt*(Xsmat(:,T-n+1)-Xhat); 
                    end
                    
                    Xsmat_hat((i-1)*3+1:(i-1)*3+3,:) = Xsmat;

                end
                
                dat.alsmsave = Xsmat_hat';
            
        end
        
    end
end