% Generate VFI Convergence Plot
% Shows how the value function converges during iteration

clear; close all;

% 1. Parameters and grid
alpha = 0.4; 
beta = 0.95; 
delta = 0.1;
kmin = 0.1; 
kmax = 10; 
N = 100;
kgrid = linspace(kmin,kmax,N)';

% 2. Utility matrix
f = kgrid.^alpha;
cons = f + (1-delta)*kgrid - kgrid';  % matrix c_{ij}
util = log(cons);                     % u(c)
util(cons<=0) = -inf;                 % infeasible choices

% 3. Value function iteration with storage
V = zeros(N,1); 
diff = 1; 
iter = 0;
V_history = [];  % store selected iterations
iter_store = [0, 5, 10, 20, 30, 50, 100];  % iterations to store
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
        [Vnew(i), ~] = max(util(i,:) + beta*V');
    end
    diff = max(abs(Vnew - V));
    V = Vnew;
    iter = iter + 1;
    
    % Save selected iterations
    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('Converged in %d iterations\n', iter);

% 4. Plot convergence
figure('Position', [100, 100, 800, 500]);

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

% Right panel: Convergence speed
subplot(1,2,2)
% Re-run to track convergence metric
V = zeros(N,1); 
diff_history = [];
iter = 0;

while iter < 1000
    Vnew = zeros(N,1);
    for i = 1:N
        [Vnew(i), ~] = max(util(i,:) + beta*V');
    end
    diff = max(abs(Vnew - V));
    diff_history(end+1) = diff;
    V = Vnew;
    iter = iter + 1;
    
    if diff < 1e-6
        break;
    end
end

semilogy(1:length(diff_history), diff_history, 'b-', 'LineWidth', 2);
grid on;
xlabel('Iteration', 'FontSize', 12);
ylabel('max|V_{t+1} - V_t|', 'FontSize', 12);
title('Convergence Metric', 'FontSize', 13);
hold on;
yline(1e-6, 'r--', 'Tolerance', 'LineWidth', 1.5, 'FontSize', 10);

% Save figure
exportgraphics(gcf,'Figures/vfi_convergence.png','Resolution',300);

