% -----------------------------------------------------------------------
% Example SSM: Solution to a Neoclassical Growth Model via Uhlig's Toolkit.
%
% Obj: Recover SMOOTHED states + build likelihood with Matlab's KS code
% using calibrated parameters, PCE data and assuming NO measurement error.
%
% The State Space Model:
% y_t = C x_t               : Observation/Measurement equation (D_t = 0)
% x_t = A x_{t-1} + B w_t   : State/ Transition equation
%
% where:
% y_t = (PCE_t) 
% x_t = (k_t, c_t, r_t, z_t)
%
% Explicitely: 1 observable and 4 latent states as AR(1).
% -----------------------------------------------------------------------

%% Houskeeping:
clear
close all
clc


%% Obtain and prepare data:
% Import real PCE data from FRED.
c = fred;
seriesID = 'PCECC96'; 
startdate = '1960-01-01';
enddate = '2019-12-31';
data = fetch(c, seriesID, startdate, enddate);
dates = data.Data(:, 1);
pce = data.Data(:, 2);
dates = datetime(dates, 'ConvertFrom', 'datenum');

% Apply the HP filter
lambda =  1600;
[trend, cycle] = hpfilter(100*log(pce), Smoothing = lambda);
log_deviation = cycle;

% Define colors and figure format
blue    = [  1,  87, 155]/255;
red     = [4/5, 1/5, 1/5];
green   = [2/5, 4/5, 2/5];
lgblue  = [108, 172, 228]/255;
violet  = [143,   0, 255]/255;
plot_colors = {blue; red; green; violet};
fig_fmt = 'png';


%% Initialize state space model:

% Matrix Parameters:
psi = 0.95; 
sigma = 0.712; 

a11 = 0.965;
a12 = 0;
a13 = 0;
a14 = 0.075*psi;
a21 = 0.618;
a22 = 0;
a23 = 0;
a24 = 0.305*psi;
a31 = -0.022;
a32 = 0;
a33 = 0;
a34 = 0.035*psi;
a41 = 0;
a42 = 0;
a43 = 0;
a44 = psi;
b11 = 0.075;
b22 = 0.305;
b33 = 0.035;
b44 = 1;
%b11 = 0.075*sigma;
%b22 = 0.305*sigma;
%b33 = 0.035*sigma;
%b44 = 1*sigma;
c11 = 0;
c12 = 1;
c13 = 0;
c14 = 0;
d11 = 0;
params_0 = [a11, a12, a13, a14, a21, a22, a23, a24, a31, a32, a33, a34, ...
            a41, a42, a43, a44, b11, b22, b33, b44, c11, c12, c13, c14, ...
            d11];


%% Initial guess for model's state equation variance-covariance matrix:
% I believe Dynare solves for the unconditional variances as:
% P = APA'+ Q; Matlab: P = dlyap(A, Q);
cov_0 =[
        17.5125, 13.3546, -0.0930, 8.3391;
        13.3546, 10.6290, -0.0211, 7.7191;
        -0.0930, -0.0211, 0.0088, 0.1497;
        8.3391, 7.7191, 0.1497, 10.2564];

%cov_0 = eye(4);

%cov_0 =[
  %    8.8779, 6.7700, -0.0471, 4.2275;
  %    6.7700, 5.3883, -0.0107, 3.9131;
  %   -0.0471, -0.0107,  0.0045, 0.0759;
  %    4.2275, 3.9131,  0.0759, 5.1994];


%% Initial estimate of the state's conditional mean (Matlab calls it as mean):
mean_0 = [0;0;0;0];  

% Define StateType: (0 for stationary)
stateType = [0, 0, 0, 0];


%% Create SSM and Filter data
% Create the state-space model by passing the function ParamMap as a function handle to ssm:
mdl = ssm(@(params) ParamMap(params, mean_0, cov_0, stateType));


% Run the filter at set parameter values to obtain filtered states:
[X, logL, SmoothedStates] = smooth(mdl, log_deviation, 'params',params_0);

% Store period log_likelihoods
log_likelihoods_seq = [SmoothedStates.LogLikelihood];

% Cumulative LL:
total_log_likelihood = sum(log_likelihoods_seq);


%% Plot Smoothed States.

long_vector = cell2mat({SmoothedStates.SmoothedStates}'); %960 in length
T = length(log_deviation);  % or 240
X = reshape(long_vector, 4, T)';  % Reshape into 4 rows, then transpose.

% 2. Align dates:
dates = dates(1:T);

% 3. Set colors:
plot_colors = {blue; red; green; violet};

% 4. State names:
state_names = {'k_t', 'c_t', 'r_t', 'z_t'};

% 5. Create single figure:
figure;
hold on;
for i = 1:4
    plot(dates, X(:,i), 'LineWidth', 1.5, 'Color', plot_colors{i});
end
hold off;
grid on;
xlabel('Year', 'FontSize', 12);
ylabel('State Value', 'FontSize', 12);
title('Smoothed States using Matlab''s KS', 'FontSize', 12);
%ylim([-0.1 0.1]); 
legend(state_names, 'Location', 'best', 'FontSize', 12);

% 6. Save figure as PNG
filename = '../eg_ssm_smoothed_2_states.png';
saveas(gcf, filename);


%% Save output
% Save the filtered states into a .mat file
save('../example_ssm.mat', 'X', 'dates');





