%% WEEK7_CHALLENGE_SOLUTION.M
% Computational Tools for Macroeconomics — Week 7
% In-class Challenge (Solution): Solow model with population growth n > 0
%
% Model per worker:
%   y_t = k_t^alpha
%   k_{t+1} = (s * k_t^alpha + (1 - delta) * k_t) / (1 + n)
%
% Steady state:
%   k* = ( s / (delta + n) )^( 1 / (1 - alpha) )
%
% Tasks solved:
%  1) Simulate k_t and y_t for n in {0, 0.01, 0.03}
%  2) Plot transition paths and steady states
%  3) Brief printed interpretation + steady-state table

clear; close all; clc;

%% --------------------------
%  Parameters (adjust freely)
%  --------------------------
alpha  = 0.4;      % capital share
s      = 0.3;      % savings rate
delta  = 0.1;      % depreciation rate
n_list = [0.00 0.01 0.03];  % population growth rates

T   = 60;          % simulation horizon (periods)
% Choose an initial k0 common to all cases.
% Start below the n=0 steady state to show monotonic convergence.
kstar_n0 = (s/(delta+0))^(1/(1-alpha));
k0  = 0.5 * kstar_n0;  % initial capital per worker

%% --------------------------
%  Containers for results
%  --------------------------
K = zeros(T, numel(n_list));
Y = zeros(T, numel(n_list));
K(1, :) = k0;                 % same initial k0 across scenarios

kstar = zeros(1, numel(n_list));
ystar = zeros(1, numel(n_list));

%% --------------------------
%  Simulate for each n
%  --------------------------
for j = 1:numel(n_list)
    n = n_list(j);

    % Steady states (per worker)
    kstar(j) = ( s / (delta + n) )^( 1 / (1 - alpha) );
    ystar(j) = kstar(j)^alpha;

    % Forward iteration
    for t = 1:T-1
        K(t+1, j) = ( s * K(t, j)^alpha + (1 - delta) * K(t, j) ) / (1 + n);
    end
    Y(:, j) = K(:, j).^alpha;
end

%% --------------------------
%  Plot: capital per worker
%  --------------------------
figure('Name','Transition paths: capital','Color','w');
lw = 1.8;

% Plot paths
hold on;
for j = 1:numel(n_list)
    plot(0:T-1, K(:, j), 'LineWidth', lw);
end

% Overlay steady states as dashed lines
colors = get(gca,'ColorOrder');
for j = 1:numel(n_list)
    yline(kstar(j), '--', sprintf('k^*(n=%.2f)', n_list(j)), ...
        'LineWidth', 1.0, 'LabelHorizontalAlignment','left');
end

grid on; box on;
xlabel('Time');
ylabel('Capital per worker, k_t');
title('Capital Transition Paths with Population Growth');
legend(arrayfun(@(x)sprintf('n=%.2f',x), n_list, 'UniformOutput', false), ...
       'Location','southeast');
set(gca,'FontSize',12);

exportgraphics(gcf, 'Figures/week7_challenge_k_paths.png', 'Resolution', 300);

%% --------------------------
%  Plot: output per worker
%  --------------------------
figure('Name','Transition paths: output','Color','w');
hold on;
for j = 1:numel(n_list)
    plot(0:T-1, Y(:, j), 'LineWidth', lw);
end

for j = 1:numel(n_list)
    yline(ystar(j), '--', sprintf('y^*(n=%.2f)', n_list(j)), ...
        'LineWidth', 1.0, 'LabelHorizontalAlignment','left');
end

grid on; box on;
xlabel('Time');
ylabel('Output per worker, y_t');
title('Output Transition Paths with Population Growth');
legend(arrayfun(@(x)sprintf('n=%.2f',x), n_list, 'UniformOutput', false), ...
       'Location','southeast');
set(gca,'FontSize',12);

exportgraphics(gcf, 'Figures/week7_challenge_y_paths.png', 'Resolution', 300);

%% --------------------------
%  Summary table (Command Window)
%  --------------------------
fprintf('\n=== Steady States (per worker) ===\n');
fprintf('Parameters: alpha=%.2f, s=%.2f, delta=%.2f\n', alpha, s, delta);
fprintf('Initial k0 = %.4f (common across cases)\n\n', k0);
fprintf('%8s  %12s  %12s\n', 'n', 'k_star', 'y_star');
for j = 1:numel(n_list)
    fprintf('%8.2f  %12.6f  %12.6f\n', n_list(j), kstar(j), ystar(j));
end
fprintf('\nInterpretation:\n');
fprintf('- Higher n (population growth) dilutes capital per worker, lowering k* and y*.\n');
fprintf('- With the same savings and depreciation, economies with higher n need more saving to\n');
fprintf('  maintain the same capital per worker.\n\n');

%% --------------------------
%  (Optional) helper function
%  --------------------------
% This function can be reused if you want to modularize the simulation further.
%{
function [k_path, y_path, k_star, y_star] = simulate_growth_with_n(alpha, s, delta, n, k0, T)
    k_star = ( s / (delta + n) )^( 1 / (1 - alpha) );
    y_star = k_star^alpha;
    k_path = zeros(T,1);
    k_path(1) = k0;
    for t = 1:T-1
        k_path(t+1) = ( s * k_path(t)^alpha + (1 - delta) * k_path(t) ) / (1 + n);
    end
    y_path = k_path.^alpha;
end
%}