%% WEEK 8 — In-Class Challenge: Solution (builds on provided code/calibration)
% Model:
%   f(k) = k^alpha, u(c) = log(c)
% Bellman (deterministic):
%   V(k) = max_{k'} { log( f(k) + (1-delta)k - k' ) + beta V(k') }
% Simulation:
%   k_{t+1} = k'(k_t),  c_t = f(k_t) + (1-delta)k_t - k_{t+1}

clear; close all; clc;
if ~exist('Figures','dir'), mkdir Figures; end

%% ---------------- 1) PARAMETERS & GRID (keep your calibration) -----------
alpha = 0.4;
beta  = 0.95;
delta = 0.1;

kmin = 0.1;
kmax = 10;
N    = 100;
kgrid = linspace(kmin, kmax, N)';

% Production and consumption matrix for this calibration
f = kgrid.^alpha;
cons = f + (1 - delta)*kgrid - kgrid.'; % c_{ij} = f(k_i) + (1-d)k_i - k'_j
util = log(cons);
util(cons <= 0) = -inf;

%% ---------------- 2) VFI with history (your convergence plot) ------------
V = zeros(N,1);
diff = 1; 
iter = 0;
V_history = [];
iter_store = [0, 5, 10, 20, 30, 50, 100];
iter_labels = {};

% Store initial guess
V_history(:,1) = V;
iter_labels{1} = 'Iter 0';

while diff > 1e-6 && iter < 1000
    Vnew = zeros(N,1);
    for i = 1:N
        % pointwise maximization
        Vnew(i) = max(util(i,:) + beta*V');
    end
    diff = max(abs(Vnew - V));
    V = Vnew;
    iter = iter + 1;

    if ismember(iter, iter_store)
        V_history(:,end+1) = V;
        iter_labels{end+1} = sprintf('Iter %d', iter);
    end
end

% Store final iteration
V_history(:,end+1) = V;
iter_labels{end+1} = sprintf('Final (Iter %d)', iter);

fprintf('VFI converged in %d iterations (baseline).\n', iter);

% ---- Plot convergence
fig_conv = figure('Position', [100, 100, 900, 480], 'Color','w');

% Left: value function evolution
subplot(1,2,1); hold on;
cmap = jet(size(V_history,2));
for j = 1:size(V_history,2)
    if j == 1
        plot(kgrid, V_history(:,j), 'k--', 'LineWidth', 1.4);
    elseif j == size(V_history,2)
        plot(kgrid, V_history(:,j), 'r-', 'LineWidth', 2);
    else
        plot(kgrid, V_history(:,j), 'Color', cmap(j,:), 'LineWidth', 1);
    end
end
grid on; xlabel('k'); ylabel('V(k)'); title('Value Function Convergence');
legend(iter_labels, 'Location','SouthEast');

% Right: convergence metric
subplot(1,2,2);
% Re-run quickly to track sup-norm
Vtmp = zeros(N,1); diff_history = []; iter2 = 0;
while iter2 < 1000
    Vnew = zeros(N,1);
    for i = 1:N
        Vnew(i) = max(util(i,:) + beta*Vtmp');
    end
    dd = max(abs(Vnew - Vtmp));
    diff_history(end+1) = dd; %#ok<SAGROW>
    Vtmp = Vnew; iter2 = iter2 + 1;
    if dd < 1e-6, break; end
end
semilogy(1:length(diff_history), diff_history, 'b-', 'LineWidth', 2); hold on;
yline(1e-6, 'r--', 'Tolerance', 'LineWidth', 1.2);
grid on; xlabel('Iteration'); ylabel('max|V_{t+1}-V_t|'); title('Convergence Metric');

exportgraphics(fig_conv, 'Figures/vfi_convergence.png', 'Resolution', 300);

%% ---------------- 3) EXTRACT POLICY (argmax) & PLOT vs 45° ---------------
% With converged V, compute argmax to get policy indices/levels
[~, pol_ind] = max(util + beta*V', [], 2);
k_policy = kgrid(pol_ind);

fig_pol = figure('Color','w'); 
plot(kgrid, k_policy, 'LineWidth', 1.5); hold on; 
plot(kgrid, kgrid, 'k--');
grid on; xlabel('k'); ylabel('k''(k)');
title(sprintf('Policy Function (\\alpha=%.2f, \\beta=%.2f, \\delta=%.2f)', alpha, beta, delta));
legend('Policy','45° line','Location','NorthWest');
exportgraphics(fig_pol, 'Figures/policy_baseline.png', 'Resolution', 300);

%% ---------------- 4) SIMULATE CAPITAL & CONSUMPTION PATH -----------------
T = 50; 
k0 = 0.5;                        % given in challenge
kpath = zeros(T,1);
kpath(1) = k0;

for t = 1:T-1
    % Interpolate policy; clip to grid if needed
    k_next = interp1(kgrid, k_policy, kpath(t), 'linear', 'extrap');
    kpath(t+1) = min(max(k_next, kmin), kmax);
end

% Consumption (respecting delta in resource constraint)
cpath = kpath(1:end-1).^alpha + (1 - delta)*kpath(1:end-1) - kpath(2:end);
if any(cpath <= 0)
    warning('Some simulated consumption values are non-positive (check grid/calibration).');
end

% Approx steady state from fixed point of policy
[~, idx_star] = min(abs(k_policy - kgrid));
kstar = kgrid(idx_star);

fig_tr = figure('Color','w'); 
subplot(2,1,1)
plot(0:T-1, kpath, 'LineWidth', 1.5); hold on; yline(kstar, ':');
grid on; xlabel('t'); ylabel('k_t'); 
title('Capital Transition (baseline)'); 
legend('k_t','k^*','Location','SouthEast');

subplot(2,1,2)
plot(0:T-2, cpath, 'LineWidth', 1.5);
grid on; xlabel('t'); ylabel('c_t'); 
title('Consumption Transition (baseline)');

exportgraphics(fig_tr, 'Figures/paths_baseline.png', 'Resolution', 300);

fprintf('Baseline approx steady state k*: %.4f\n', kstar);

%% ---------------- 5) PARAMETER EXPERIMENTS (same grid & delta) ----------
beta_set  = [0.90, 0.95, 0.98];
alpha_set = [0.25, 0.30, 0.40];

fig_pols = figure('Color','w'); hold on;
fig_kpaths = figure('Color','w'); hold on;
legtxt = {};

for ib = 1:numel(beta_set)
    for ia = 1:numel(alpha_set)
        beta_i  = beta_set(ib);
        alpha_i = alpha_set(ia);

        % Rebuild util matrix for (alpha_i, beta_i) with SAME delta & SAME GRID
        f_i = kgrid.^alpha_i;
        cons_i = f_i + (1 - delta)*kgrid - kgrid.';
        U_i = log(cons_i); U_i(cons_i <= 0) = -inf;

        % VFI (quick, no history needed here)
        V_i = zeros(N,1); 
        for it = 1:1000
            Vnew_i = max(U_i + beta_i*V_i', [], 2);
            if max(abs(Vnew_i - V_i)) < 1e-6
                V_i = Vnew_i; break;
            end
            V_i = Vnew_i;
        end
        [~, pol_i] = max(U_i + beta_i*V_i', [], 2);
        kpol_i = kgrid(pol_i);

        % Plot policy
        figure(fig_pols);
        plot(kgrid, kpol_i, 'LineWidth', 1.1);
        legtxt{end+1} = sprintf('\\alpha=%.2f, \\beta=%.2f', alpha_i, beta_i); %#ok<SAGROW>

        % Simulate from same k0 for T
        ksim = zeros(T,1); ksim(1) = k0;
        for t = 1:T-1
            k_next = interp1(kgrid, kpol_i, ksim(t), 'linear', 'extrap');
            ksim(t+1) = min(max(k_next, kmin), kmax);
        end

        % Plot capital paths
        figure(fig_kpaths);
        plot(0:T-1, ksim, 'LineWidth', 1.1);
    end
end

% Finalize figures
figure(fig_pols);
plot(kgrid, kgrid, 'k--'); 
grid on; xlabel('k'); ylabel('k''(k)');
title(sprintf('Policies across calibrations (\\delta=%.2f, same grid)', delta));
legend(legtxt{:}, '45° line', 'Location','NorthWest');
exportgraphics(fig_pols, 'Figures/policies_grid.png', 'Resolution', 300);

figure(fig_kpaths);
grid on; xlabel('t'); ylabel('k_t');
title(sprintf('Capital paths across calibrations (k_0=%.2f)', k0));
legend(legtxt{:}, 'Location','SouthEast');
exportgraphics(fig_kpaths, 'Figures/paths_grid.png', 'Resolution', 300);

%% ---------------- 6) QUICK INTERPRETATION LINES (console) ---------------
fprintf('\nINTERPRETATION NOTES:\n');
fprintf('- Higher beta (more patience) shifts policy up, raising k^* and the k_t path.\n');
fprintf('- Higher alpha (higher MPK) increases desired saving, raising k^* and speeding convergence.\n');
fprintf('- With the same delta and grid, differences isolate preferences/technology effects.\n');