%% Coleman's Policy Function Iteration
% Solves the stochastic neoclassical growth model using time iteration on the Euler equation.

clear; close all; clc;

%% 1. Calibration
alpha = 0.33;       % Capital share
beta  = 0.96;       % Discount factor
delta = 0.10;       % Depreciation rate
sigma = 2;          % Coefficient of relative risk aversion (utility curvature)
rho   = 0.90;       % Persistence of productivity shock
sigma_eps = 0.02;   % Std dev of shock innovation

% Utility function and its derivative
u  = @(c) (c.^(1-sigma) - 1) / (1-sigma);
du = @(c) c.^(-sigma);

%% 2. Grids
% Productivity (Tauchen)
Nz = 5;             % Number of productivity states
m  = 3;             % Number of std devs
[Z, P] = tauchen(rho, sigma_eps, Nz, m);
% Z is log productivity, so A = exp(Z)
A_grid = exp(Z);

% Capital
Nk = 100;
k_min = 0.1;
k_max = 15.0;
K_grid = linspace(k_min, k_max, Nk)';

%% 3. Initialization
% Initial guess for the policy function k'(k, z)
% Guess: consume fraction of resources, or stay constant
% A simple guess is k'(k) = k (steady state-ish)
K_policy_old = repmat(K_grid, 1, Nz); 
K_policy_new = zeros(Nk, Nz);

max_iter = 1000;
tol = 1e-5;
diff = 1;

disp('Starting Coleman Policy Function Iteration...');
tic;

%% 4. Iteration
for iter = 1:max_iter
    
    for iz = 1:Nz
        A = A_grid(iz);
        
        for ik = 1:Nk
            k_curr = K_grid(ik);
            
            % Total resources available today
            % y = A * k^alpha
            resource = A * k_curr^alpha + (1-delta)*k_curr;
            
            % Define the residual function for the Euler equation
            % We search for k_next in [k_min, resource - small_eps]
            
            euler_resid = @(k_next) solve_euler(k_next, resource, ...
                K_policy_old, Z, P, iz, beta, alpha, delta, du, A_grid, K_grid);
            
            % Find k_next that zeroes the Euler equation
            % Bounds: k_next must be positive and leave some consumption
            lb = k_min;
            ub = max(k_min, resource - 1e-6); 
            
            if euler_resid(lb) * euler_resid(ub) > 0
                % If signs are same, solution might be at boundary
                if abs(euler_resid(lb)) < abs(euler_resid(ub))
                    k_opt = lb;
                else
                    k_opt = ub;
                end
            else
                try
                    k_opt = fzero(euler_resid, [lb, ub]);
                catch
                    % Fallback if fzero fails (rare in this setup)
                    k_opt = (lb+ub)/2;
                end
            end
            
            K_policy_new(ik, iz) = k_opt;
        end
    end
    
    % Check convergence
    diff = max(abs(K_policy_new(:) - K_policy_old(:)));
    if mod(iter, 10) == 0
        fprintf('Iteration %d, Error: %.2e\n', iter, diff);
    end
    
    if diff < tol
        fprintf('Converged in %d iterations.\n', iter);
        break;
    end
    
    % Update guess
    K_policy_old = K_policy_new;
    
end

time_coleman = toc;
fprintf('Elapsed time (Coleman): %.4f seconds\n', time_coleman);

%% 5. Plotting
figure;
plot(K_grid, K_policy_new, 'LineWidth', 2);
hold on;
plot(K_grid, K_grid, 'k--', 'LineWidth', 1);
xlabel('Capital k');
ylabel('Next Capital k''');
title('Policy Function (Coleman Method)');
% Construct legend labels correctly
legend_labels = cell(length(A_grid) + 1, 1);
for i = 1:length(A_grid)
    legend_labels{i} = sprintf('A = %.2f', A_grid(i));
end
legend_labels{end} = '45-degree line';
legend(legend_labels, 'Location', 'best');
grid on;
exportgraphics(gcf, 'Figures/coleman_policy.png', 'Resolution', 300);

% Save the policy for simulation AND the execution time
save('policy_coleman.mat', 'K_grid', 'Z', 'A_grid', 'K_policy_new', 'P', 'alpha', 'beta', 'delta', 'sigma', 'time_coleman');


%% Local Function for Euler Residual
function resid = solve_euler(k_next, resource, K_policy_future, Z_grid, P, iz_curr, beta, alpha, delta, du, A_grid, K_grid)
    % Current consumption implied by k_next
    c_curr = resource - k_next;
    
    if c_curr <= 0
        resid = 1e10; % Penalty
        return;
    end
    
    % LHS of Euler: u'(c_t)
    mu_curr = du(c_curr);
    
    % RHS: beta * E [ u'(c_{t+1}) * (f'(k_{t+1}) + 1 - delta) ]
    expectation_term = 0;
    
    % Integrate over future productivity states z'
    for iz_next = 1:length(Z_grid)
        prob = P(iz_curr, iz_next);
        A_next = A_grid(iz_next);
        
        % Interpolate future policy k'' = g(k', z')
        k_next2 = interp1(K_grid, K_policy_future(:, iz_next), k_next, 'linear', 'extrap');
        
        % Future consumption
        resource_next = A_next * k_next^alpha + (1-delta)*k_next;
        c_next = resource_next - k_next2;
        
        if c_next <= 0
            c_next = 1e-6; % Avoid complex numbers
        end
        
        % Future marginal product of capital (MPK)
        mpk_next = alpha * A_next * k_next^(alpha-1);
        return_next = mpk_next + 1 - delta;
        
        expectation_term = expectation_term + prob * du(c_next) * return_next;
    end
    
    resid = mu_curr - beta * expectation_term;
end
