%% week6_challenge_starter.m
% Week 6 — In-class Challenge: Calibrate Solow savings rate s
% Goal: Find s that matches an observed output path y_data.
% Steps: (i) manual grid search; (ii) (bonus) fminsearch with reparam.

clear; close all; clc;

% Ensure Figures folder exists
if ~exist('Figures','dir'), mkdir('Figures'); end

%% ------------------ Model & Data Setup ------------------
% Parameters (fixed)
params.alpha = 0.33;   % capital share
params.delta = 0.08;   % depreciation
params.n     = 0.01;   % population growth
params.g     = 0.02;   % technology growth
params.A0    = 1.0;    % initial technology
params.K0    = 1.0;    % initial capital
T            = 40;     % horizon

% Synthetic "observed" data generated with s_true + small noise
s_true   = 0.24;
y_model0 = solow_simulate(s_true, params, T);  % true model path
rng(123);
y_data   = y_model0 .* exp(0.05*randn(size(y_model0))); % ~noise-free for class

% Objective: SSE between model and data for a given s
obj = @(s) sum((solow_simulate(s, params, T) - y_data).^2);

%% ------------------ (1) Manual Grid Search ------------------
s_grid = linspace(0.05, 0.5, 200);
SSE    = zeros(size(s_grid));
for i = 1:length(s_grid)
    SSE(i) = obj(s_grid(i));
end
[~, idx_min] = min(SSE);
s_hat_grid   = s_grid(idx_min);

fprintf('Grid search estimate: s_hat = %.4f (true s = %.4f)\n', s_hat_grid, s_true);

% Plot objective over grid
fig1 = figure('Color','w','Position',[100 100 720 480]);
plot(s_grid, SSE, 'LineWidth', 1.8); hold on;
plot(s_hat_grid, SSE(idx_min), 'o', 'MarkerSize',8, 'LineWidth',1.5);
xlabel('Savings rate s'); ylabel('SSE');
title('Objective: sum of squared errors vs s (Grid Search)');
grid on; box on;
print(fig1, fullfile('Figures','fig_solow_obj_grid.pdf'), '-dpdf', '-r300');
print(fig1, fullfile('Figures','fig_solow_obj_grid.png'), '-dpng', '-r300');

%% ------------------ (2) BONUS: fminsearch with reparam ------------------
% Re-parameterize s = sigmoid(x) to enforce 0 < s < 1
sigmoid = @(x) 1./(1 + exp(-x));
obj_x   = @(x) obj(sigmoid(x));

x0 = 0;  % corresponds to s0 = 0.5
[x_hat, fval] = fminsearch(obj_x, x0);
s_hat_nelder  = sigmoid(x_hat);

fprintf('fminsearch estimate (reparam): s_hat = %.4f, SSE = %.4f\n', s_hat_nelder, fval);

% Compare fitted vs data
y_fit = solow_simulate(s_hat_nelder, params, T);
fig2 = figure('Color','w','Position',[100 100 720 480]);
plot(1:T, y_data, '-','LineWidth',1.8); hold on;
plot(1:T, y_fit, '--','LineWidth',1.8);
xlabel('t'); ylabel('Output y_t');
title('Solow Calibration Fit: Data vs Model');
legend('Data','Model (ŝ)','Location','best'); grid on; box on;
print(fig2, fullfile('Figures','fig_solow_fit_inclass.pdf'), '-dpdf', '-r300');
print(fig2, fullfile('Figures','fig_solow_fit_inclass.png'), '-dpng', '-r300');

%% ------------------ Helper: Solow simulator ------------------
function y = solow_simulate(s, P, T)
    % Clamp s to (0,1) for numerical safety in class
    s = max(min(s, 0.999), 0.001);
    K = zeros(T,1); A = zeros(T,1); Y = zeros(T,1);
    K(1) = P.K0; A(1) = P.A0;
    for t = 1:T
        Y(t) = K(t)^P.alpha * A(t)^(1-P.alpha);         % Y_t = K^alpha (A)^(1-alpha), L normalized
        if t < T
            K(t+1) = (1 - P.delta)*K(t) + s*Y(t);
            A(t+1) = (1 + P.g)*A(t);
        end
    end
    y = Y;  % observed output path (per effective worker scaling absorbed by A)
end