%% Parse .hmod Files
% Suppose you have KS_JEDC10.hmod in the working directory. Call in MATLAB

clear_hans
hans KS_JEDC10

% As shown, HANS parases the hmod file, compiles the mex files, and returns a structure that describes the model.
% clear_hans in the first line clears any cached files generated before.
%
%% Solve the Stationary Equilibrium
% Call the generated code solve_ss to solve the stationary equilibrium

tic;
ss_rslt = solve_ss;
fprintf('Time used for solving the stationary equilibrium:\n');toc;
ss_rslt

% As shown, the function returns a structure that describes the solved stationary equilibrium.
%
% We can inspect the values of equilibrium variables:

ss_rslt.var_agg

% We can insepct the stationary distribution

plot(ss_rslt.vfi_rslt.var_state.a, ss_rslt.dist);
xlabel('k');
ylabel('fractions');

% and the policy functions

plot(ss_rslt.vfi_rslt.var_state.a, ss_rslt.vfi_rslt.var_policy.ap);
xlabel('a');
title('Policy functions for ap');
legend({'low e','high e'})

%
%% Solve the Transition Path After A Temporay Shock
%% Solve the linearized transition path
% We first construct the sequence of productivity shock.

T = 200;
shock_Z = 0.01; rho_Z = 0.9;
Z_t = 1 + shock_Z *rho_Z.^(0:(T-1));
figure;
plot(Z_t)

% We then pass the productivity shock and the solved ss_rslt into solve_trans_linear to solve the linearized transition path.

options = struct;
options.T = T;
options.Z_t = Z_t;
tic;
trans_linear_rslt = solve_trans_linear(ss_rslt, options);
fprintf('Time used for solving the linearized transition path:\n');toc;
trans_linear_rslt

% We can plot the transition path (note that K is the next period capital).

figure;
plot(trans_linear_rslt.var_agg_t.K)
title('K');
xlabel('time');

%
%% Solve the nonlinear transition path
% Now we pass the solved ss_rslt to solve_nonlinear_trans(init_ss, final_ss, options) to solve the nonlinear transition path.
% Note that we pass ss_rslt as both init_ss and final_ss.

tic;
trans_nonlinear_rslt = solve_trans_nonlinear(ss_rslt, ss_rslt, options);
fprintf('Time used for solving the nonlinear transition path:\n');toc;
trans_nonlinear_rslt

% We compare the solution of the nonlinear transition path and the linear transition path.

figure;
plot(trans_linear_rslt.var_agg_t.K);
hold on;
plot(trans_nonlinear_rslt.var_agg_t.K, '--');
xlabel('time');
title('K');
legend({'Linear','Nonlinear'});

% We can also inspect other defined equilibrium varialbes along the transition path.

figure;
plot(trans_nonlinear_rslt.var_agg_t.w);
title('w');
xlabel('time');

%
%% Solve the Transition Path After A Permanent Shock
% We now consider an unexpected shock that raises Z permanently by 100%.
% We first sovle the new stationary equilibrium.

options = struct;
options.Z = 2.0;
ss_rslt_new = solve_ss(options);

% We then solve the transition path from the initial stationary equilibrium to the new stationary equilibrium.

options = struct;
T = 200;
options.Z_t = 2.0*ones(1,T);  % shock hits at period 1
options.T = T;
tic;
trans_to_new_ss_rslt = solve_trans_nonlinear(ss_rslt, ss_rslt_new, options);
fprintf('Time for solving the nonlinear transition path after a permanent shock:\n');toc;
trans_to_new_ss_rslt

% We next inspect the nonlinear transition path.

figure;
subplot(2,2,1);
plot(trans_to_new_ss_rslt.var_agg_t.K); title('K');
subplot(2,2,2);
plot(trans_to_new_ss_rslt.var_agg_t.C); title('C');
subplot(2,2,3);
plot(trans_to_new_ss_rslt.var_agg_t.w); title('w');
subplot(2,2,4);
plot(trans_to_new_ss_rslt.var_agg_t.r); title('r');

%
%% Simulate Individual Samples
%% Simulate Individual Samples at the Stationary Distribution
% We can simulate individual samples at the stationary equilibrium or along the transition path.
% Let's start with the stationary equilibrium.

tic;
simu_rslt = simulate_stochastic(ss_rslt);
fprintf('Time used for simulating individual samples at the stationary equilibrium:\n');toc;
simu_rslt

% As shown, passing ss_rslt to simulate_stochastic returns a panel of simualted individual samples.
% Printed information is the mean of variables at the printed period.
% By default, it simulates 10,000 samples for 1,000 periods, and starts from random state values drawn from the stationary distribution.
% We can overwrite the num_samples and num_periods:

simu_options = struct;
simu_options.num_samples = 2e4;
simu_options.num_periods = 2e3;
tic;
simu_rslt = simulate_stochastic(ss_rslt, simu_options);
fprintf('Time used for simulating 2e4 individual samples for 2e3 periods at the stationary equilibrium:\n');toc;

% We can compare the histogram of simulated individual samples to the one constructed based on the non-stochastic simulation.

figure; 
subplot(1,2,1);
histogram(simu_rslt.a(:,end),ss_rslt.vfi_rslt.var_state.a,'Normalization','probability');
xlim([-10,210]);ylim([0,0.014]);
subplot(1,2,2);
plot(ss_rslt.vfi_rslt.var_state.a, sum(ss_rslt.dist,1)');
xlim([-10,210]);ylim([0,0.014]);

%
%% Simulate Individual Samples Along the Transition Path
% To simulate individual samples along the transition path, we need to specify values of initial state values.
% Let's specify them to be drawn from the stationary distribution.

simulate_initials = struct;
% Use the last period of the simulated results at the stationary distirubtion
simulate_initials.shock = simu_rslt.shock(:,end);
simulate_initials.a = simu_rslt.a(:,end);

% We now pass the initial values into options, and simulate individual samples along the transition path to the new stationary equilibrium.

simu_options = struct;
simu_options.num_samples = 2e4;  % to be consistent with size of initial values
simu_options.simulate_initials = simulate_initials;
simu_trans_rslt = simulate_stochastic(trans_to_new_ss_rslt, simu_options);

% We then compare the aggregate variables constructed from individual samples and from non-stochastic simulations.

figure;
plot(trans_to_new_ss_rslt.var_agg_t.c,'-');
hold on;
plot(mean(simu_trans_rslt.c),'--');
legend({'Determinstic', 'Stochastic'},'Location','Best');
title('c');
xlabel('t');

%