%% Endogenous Grid Method (EGM)
% Solves the stochastic neoclassical growth model using EGM.
% Comparison with Coleman's method for efficiency.

clear; close all; clc;

%% 1. Calibration
alpha = 0.33;       
beta  = 0.96;       
delta = 0.10;       
sigma = 2;          
rho   = 0.90;       
sigma_eps = 0.02;   

u_prime  = @(c) c.^(-sigma);
u_inv_prime = @(mu) mu.^(-1/sigma); % Inverse of marginal utility

%% 2. Grids
Nz = 5;             
m  = 3;             
[Z, P] = tauchen(rho, sigma_eps, Nz, m);
A_grid = exp(Z);

Nk = 100;
k_min = 0.1;
k_max = 15.0;
K_grid = linspace(k_min, k_max, Nk)'; % This acts as the grid for k' (tomorrow's capital)

%% 3. Initialization
% Policy function guess: K'(k,z) defined on the fixed grid K_grid
K_policy_old = repmat(K_grid, 1, Nz);
K_policy_new = zeros(Nk, Nz);

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

disp('Starting Endogenous Grid Method (EGM)...');
tic;

%% 4. Iteration
for iter = 1:max_iter
    
    % Pre-calculate Consumption at grid points under current policy
    % C(k, z) = Production(k, z) + (1-delta)k - K'(k, z)
    % Here "k" refers to the grid points which we treat as k' for the expectation step,
    % but we also need the policy function values at these points.
    
    Cons_old = zeros(Nk, Nz);
    for iz = 1:Nz
        A = A_grid(iz);
        Resources = A .* K_grid.^alpha + (1-delta).*K_grid;
        Cons_old(:, iz) = max(1e-6, Resources - K_policy_old(:, iz));
    end
    
    Mu_old = u_prime(Cons_old); % Marginal utility at grid points
    
    % Loop over current productivity state z
    for iz = 1:Nz
        
        % Calculate Expected Marginal Utility for each choice of k' in K_grid
        % E_mu = sum( P(iz, :) * u'(c') * R' )
        
        Expected_RHS = zeros(Nk, 1);
        
        for iz_next = 1:Nz
            prob = P(iz, iz_next);
            A_next = A_grid(iz_next);
            
            % Future Return R' = alpha * A' * (k')^(alpha-1) + 1 - delta
            % Note: here k' is simply the grid K_grid
            R_next = alpha * A_next * K_grid.^(alpha-1) + 1 - delta;
            
            % Add to expectation: prob * u'(c') * R'
            % Since K_grid is exactly the grid where Mu_old is defined, we use it directly!
            Expected_RHS = Expected_RHS + prob * Mu_old(:, iz_next) .* R_next;
        end
        
        % Euler Equation: u'(c) = beta * Expected_RHS
        % Solve for c:
        c_implied = u_inv_prime(beta * Expected_RHS);
        
        % Find endogenous k that supports this choice:
        % Resource(k) = c_implied + k'
        % A * k^alpha + (1-delta)*k = c_implied + K_grid
        
        Total_Assets_Required = c_implied + K_grid;
        
        % Invert the resource constraint to find k_endo
        % Since A k^alpha + (1-delta)k is increasing monotonic, this is well defined.
        % We can use a simple root finder or fzero loop.
        
        k_endo = zeros(Nk, 1);
        A_curr = A_grid(iz);
        
        for i = 1:Nk
            target = Total_Assets_Required(i);
            resource_func = @(k) A_curr * k^alpha + (1-delta)*k - target;
            
            % Good initial guess is helpful
            k_guess = target; % roughly
            
            try
                k_endo(i) = fzero(resource_func, [0.01, target*2]);
            catch
                % fallback
                 k_endo(i) = fzero(resource_func, 0.01);
            end
        end
        
        % Now we have pairs (k_endo, K_grid). 
        % We want policy on the fixed K_grid.
        % k_endo is the state k today. K_grid is the choice k' today.
        % Interpolate to get K'(K_grid, z)
        
        K_policy_new(:, iz) = interp1(k_endo, K_grid, K_grid, 'linear', 'extrap');
    end
    
    % Enforce boundaries
    K_policy_new(K_policy_new < k_min) = k_min;
    % No strict upper bound needed usually, but good to clamp for stability if diverging
    
    % 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
    
    K_policy_old = K_policy_new;
end

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

%% 5. Comparison and Plotting
% Load Coleman results if available
try
    coleman = load('policy_coleman.mat');
    
    % Calculate time difference if variable exists
    if isfield(coleman, 'time_coleman')
        fprintf('Time Difference: EGM is %.2fx faster than Coleman.\n', coleman.time_coleman / time_egm);
    else
        disp('Variable time_coleman not found in policy_coleman.mat');
    end

    
    figure;
    subplot(1,2,1)
    plot(K_grid, K_policy_new, 'LineWidth', 2);
    title('EGM Policy Function');
    xlabel('k'); ylabel('k''');
    grid on;
    
    subplot(1,2,2)
    % Plot difference between methods
    diff_methods = abs(K_policy_new - coleman.K_policy_new);
    plot(K_grid, diff_methods);
    title('Difference (EGM - Coleman)');
    xlabel('k'); ylabel('Absolute Difference');
    grid on;
    
    exportgraphics(gcf, 'Figures/EGM_vs_Coleman.png', 'Resolution', 300);
catch
    disp('Coleman results not found for comparison. Run coleman.m first.');
end

% Save EGM results
save('policy_egm.mat', 'K_grid', 'Z', 'A_grid', 'K_policy_new', 'P', 'alpha', 'beta', 'delta', 'sigma');

