%% VFI code

clear; close; clc;

% 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 = -inf(N);
feas = cons > 0;
util(feas) = log(cons(feas));

% 3. Value function iteration
V = zeros(N,1); diff = 1;
while diff > 1e-6
    Vnew = zeros(N,1);
    for i = 1:N
        [Vnew(i), pol_ind(i)] = max(util(i,:) + beta*V');
    end
    diff = max(abs(Vnew - V));
    V = Vnew;
end

% 4. Policy function and value function
k_policy = kgrid(pol_ind);

% 5. Plot
figure;
subplot(1,2,1)
plot(kgrid, V)
title('Value Function'); xlabel('k'); ylabel('V(k)')

subplot(1,2,2)
plot(kgrid, k_policy, 'LineWidth',1.2)
hold on; plot(kgrid, kgrid, '--k')
title('Policy Function'); xlabel('k'); ylabel('k''(k)')
legend('Policy','45° line','Location','NorthWest')
exportgraphics(gcf,'Figures/VFI.png','Resolution',300);


%%

% 6. Simulate transition dynamics
T = 50;                        % number of periods
kpath = zeros(T,1);
kpath(1) = 0.5;                % initial capital

for t = 1:T-1
    % interpolate policy to get next capital
    kpath(t+1) = interp1(kgrid, k_policy, kpath(t));
end


% Plot transition
figure;
plot(1:T, kpath, 'LineWidth',1.2)
xlabel('Time'); ylabel('Capital')
title('Transition of Capital over Time')
exportgraphics(gcf,'Figures/VFI_K_path.png','Resolution',300);
