
function [aggregates_old, aggregates_new] = solve_idiosyncratic_risk()


% Assigning parameters
params.discount_factor = 0.96                        ;
params.interest_rate   = 1.04                        ;
params.wage            = 1                           ;
params.productivity    = [0.3, 1]                    ;
params.n_productivity  = length(params.productivity) ;
params.markov_matrix   = [0.30, 0.70; 0.05, 0.95]    ;


% Discretizing the state space
asset_lower_bound = 0.001                                                    ;
asset_upper_bound = 15                                                       ;
n_assets          = 50                                                       ;
lower_bound       = log(asset_lower_bound)/log(10)                           ;
upper_bound       = log(asset_upper_bound)/log(10)                           ;
assets            = logspace( lower_bound, upper_bound, n_assets)            ;
params.n_assets   = n_assets                                                 ;
params.assets     = assets                                                   ;


aggregates_old = solve_aggregates(params);

params.productivity    = [0.15, 1]               ;
aggregates_new         = solve_aggregates(params);

saving_change = (aggregates_new.aggregate_savings - aggregates_old.aggregate_savings)/aggregates_old.aggregate_savings;
fprintf('Change in aggregate savings is %0.8f%%.\n',  saving_change*100)



end


function aggregates = solve_aggregates(params)

% Solving agent problem, then supplying agent arrays to solve stationary
% distribution
optimal_values = solve_infinite_horizon(params);
[stationary_distribution] = solve_stationary_distribution( optimal_values, params );

aggregator = @(x) sum(x(:).*stationary_distribution(:));

aggregates.aggregate_savings = aggregator( optimal_values.savings        );
aggregates.aggregate_welfare = aggregator( optimal_values.value_function );

end




function [stationary_distribution] = solve_stationary_distribution( optimal_values, params )


% Generate optimal values if not provided
if ~exist('optimal_values', 'var') || isempty(optimal_values)
    optimal_values = solve_infinite_horizon(params);
end

distribution = zeros( params.n_assets, length(params.productivity) ) ;
distribution = distribution + 1 / ( params.n_assets * length(params.productivity) ) ;    % Initial guess
distribution_next = zeros( size(distribution) );
tolerance = 10e-08;

while true    
    for id = 1:params.n_productivity
        for ia = 1:params.n_assets
            
            index_upper = find( optimal_values.savings( ia, id ) < params.assets, 1, 'first' );
            if ~isempty(index_upper)
                w_upper = (optimal_values.savings( ia, id ) - params.assets(index_upper-1))/(params.assets(index_upper) - params.assets(index_upper-1));
            elseif isempty(index_upper)
                w_upper = 1;
                index_upper = params.n_assets;
            end

            for idn = 1:params.n_productivity
                
                distribution_next( index_upper    , idn ) = distribution_next( index_upper    , idn ) +      w_upper  * distribution(ia,id)*params.markov_matrix(id,idn);
                distribution_next( index_upper - 1, idn ) = distribution_next( index_upper - 1, idn ) + (1 - w_upper) * distribution(ia,id)*params.markov_matrix(id,idn);
                
            end
        end
    end
    
    % Testing for sum of distribution close to one
    if abs( sum(distribution_next(:) ) - 1 ) > 10e-5
        error('Distribution failed to coverge.')
    end
    
    distance = max(abs(distribution(:) - distribution_next(:)));
    fprintf('The current error is %0.10f.\n', distance)
    distribution = distribution_next;
    distribution_next = zeros(size(distribution));
    if distance<tolerance
        break
    end
end

stationary_distribution = distribution;

end




function optimal_values = solve_infinite_horizon(params)

optim_options = optimset('Display', 'off', 'TolFun', 1e-5, 'TolX', 1e-5, 'MaxFunEvals', 10e4);

% Preallocating arrays
value_function      = zeros( params.n_assets, params.n_productivity ) ;    % Also initial guess of value function.
savings             = zeros( params.n_assets, params.n_productivity ) ;
new_value_function  = zeros( params.n_assets, params.n_productivity ) ;

% Solving for the value function
tolerance = 10e-5 ;
iteration = 0     ;
max_iter = 500    ;

while true
    
    iteration = iteration + 1;    
    for ia = 1:params.n_assets
        for ip = 1:params.n_productivity
            asset               = params.assets(ia)                                   ;
            initial_guess       = params.assets(ia)                                   ;
            productivity        = params.productivity(ip)                      ;
            continuation_values = value_function * params.markov_matrix(ip,:)' ;    % Takes expectation over continuation values
            [optimal_saving, optimal_value] = fminsearch(@(saving) solve_value(saving, asset, productivity, continuation_values, params),...
                                                initial_guess, optim_options)  ;
            new_value_function( ia,ip ) = -1*optimal_value                     ;
            savings           ( ia,ip ) = optimal_saving                       ;
        end
    end
    
    distance = max( abs( value_function(:) - new_value_function(:) ) )   ;
    fprintf('Agent convergence criteria is %0.4f.\n', distance)
    if ( distance<tolerance ) || ( iteration>max_iter )
        break
    end
    
    % Updating value function guess
    value_function = new_value_function;
    
end

optimal_values.value_function = value_function;
optimal_values.savings        = savings       ; 

figure
plot( params.assets, savings(:,1), params.assets, savings(:,2) ,params.assets, params.assets )
legend('Unemployed','Employed','45 degree line','Location','SouthEast')
xlabel('Assets')
ylabel('Savings')

end


function lifetime_utility = solve_value(saving, asset, productivity, continuation_values, params)

% Testing Constraint
if ( saving<params.assets(1) ) || ( saving>params.assets(end) )
    lifetime_utility = inf;
    return
end

% Solving consumption and testing for positivity
consumption = params.interest_rate*asset - saving + params.wage * productivity ;

if ( consumption < 0 )
    lifetime_utility = inf;
    return
end

% Calculating continuation value
continuation_value = interp1( params.assets, continuation_values, saving, 'linear') ;

lifetime_utility = log( consumption ) + params.discount_factor * continuation_value ;
lifetime_utility = -1*lifetime_utility;



end