function [optimal_aggregates, no_ss_aggregates] = calculate_optimal_benefit()

benefit_ub   = 2.5                                            ;
benefit_lb   = 0                                              ;
n_benefits   = 12                                             ;
benefits     = linspace( benefit_lb, benefit_ub, n_benefits ) ;
welfares     = zeros( 1, n_benefits )                         ;
benefit_taxs = zeros( 1, n_benefits )                         ;

parfor ib = 1:n_benefits
    fprintf('Solving for benefit %d.\n', ib)
    [welfares( ib ), benefit_taxs(ib)] = evaluate_financed_benefit( benefits( ib ) );
end

plot(benefits,welfares)
xlabel('Benefit')
ylabel('Welfare')

% Defining Balanced-budget Social Welfare as a function of the benefit
social_security_welfare = @(benefit) -interp1(benefits,welfares,benefit,'spline','extrap');
initial_guess = 0.5*(benefit_lb + benefit_ub);
optimal_benefit = fminsearch(@(benefit) social_security_welfare(benefit), initial_guess);

% Finding the optimal replacement ratio
optimal_benefit_tax = interp1( benefits, benefit_taxs, optimal_benefit, 'spline', 'extrap');
optimal_aggregates = solve_lifecycle_model( optimal_benefit, optimal_benefit_tax);
optimal_replacement_ratio = optimal_aggregates.benefit/optimal_aggregates.labor_income;

fprintf('The optimal replacement ratio is %0.2f percent.\n', optimal_replacement_ratio*100)

% Solving the aggregates in an economy without Social Security
no_ss_aggregates = solve_lifecycle_model( 0, 0);
consumption_change  = (no_ss_aggregates.consumption  - optimal_aggregates.consumption  )/optimal_aggregates.consumption  ;
savings_change      = (no_ss_aggregates.saving       - optimal_aggregates.saving       )/optimal_aggregates.saving       ;
labor_supply_change = (no_ss_aggregates.labor_supply - optimal_aggregates.labor_supply )/optimal_aggregates.labor_supply ;
welfare_change      = (no_ss_aggregates.welfare      - optimal_aggregates.welfare      )/optimal_aggregates.welfare      ;

fprintf( 'The change in consumption resulting from SS is %0.2f percent.\n' , consumption_change *100 )
fprintf( 'The change in savings resulting from SS is %0.2f percent.\n'     , savings_change     *100 )
fprintf( 'The change in labor supply resulting from SS is %0.2f percent.\n', labor_supply_change*100 )
fprintf( 'The change in welfare resulting from SS is %0.2f percent.\n'     , welfare_change     *100 )

end




function [welfare, benefit_tax] = evaluate_financed_benefit(benefit)


tax_ub        = 1     ;
tax_lb        = 0     ;
tolerance     = 10e-5 ;
iteration     = 0     ;
max_iteration = 50    ;

while true
    
    % Updating iteration and testing for iteration limit
    iteration = iteration + 1;
    if iteration > max_iteration
        error('Maximum number of iterations reached.')
    end
    
    % Updating tax guess, solving aggregates and budget items
    test_tax   = 0.5*(tax_lb + tax_ub)                       ;
    aggregates = solve_lifecycle_model(benefit,test_tax)     ;
    surplus    = aggregates.tax_revenue - aggregates.benefit ;
    
    % Testing convergence
    if abs( surplus ) < tolerance
        break
    end
    
    % Updating guess
    if surplus > 0
        tax_ub = test_tax ;
    elseif surplus < 0
        tax_lb = test_tax ;
    end
end

% Assigning function's outputs
welfare     = aggregates.welfare ;
benefit_tax = test_tax           ; 

end




function aggregates = solve_lifecycle_model( benefit, tax_rate )

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

% Assigning parameters
discount_factor        = 1.01 ;
crra                   = 3    ;
consumption_share      = 0.35 ;
wage                   = 1    ;
interest_rate          = 1.04 ;
initial_wealth         = 0    ;
ret_age                = 45   ;
max_age                = 80   ;
population_growth_rate = 0.012;

% Reading in productivty and survival profiles from file
productivity           = xlsread( 'productivity_profile.xls'    ); % Estimated 2012 Texas log hourly earnings ages 21-80
survival_probabilities = xlsread( 'survival_probabilities.xls' ); % Estimated survival probabilities from the SSA

%  Adjusting productivity for retirement
productivity = [productivity(1:ret_age-1)', zeros(max_age - (ret_age-1),1)']';

% Storing values in a structure
params.discount_factor        = discount_factor        ;
params.crra                   = crra                   ;
params.consumption_share      = consumption_share      ;
params.wage                   = wage                   ;
params.interest_rate          = interest_rate          ;
params.initial_wealth         = initial_wealth         ;
params.ret_age                = ret_age                ;
params.max_age                = max_age                ;
params.productivity           = productivity           ;
params.survival_probabilities = survival_probabilities ;
params.benefit                = benefit                ;
params.tax_rate               = tax_rate               ;

% Calculating population weights
population_weights    = zeros( 1, max_age ) ;
population_weights(1) = 1                   ;    % Initiating population
for age = 2:max_age
    population_weights(age) = survival_probabilities(age-1)*population_weights(age-1)...
                              /(1 + population_growth_rate)                             ;
end
population_weights = population_weights./sum(population_weights);


% Discretizing the state space
asset_upper_bound = 10  ;    % Take this as given now...we'll have more to say about it later.
asset_lower_bound =  0.1;    % A negative value here means we can borrow.
n_assets          = 80  ;    % Be careful here.  A large value will take a long time to run.
assets            = linspace( asset_lower_bound, asset_upper_bound, n_assets);

% Loading 'assets' in the 'params' structure
params.assets = assets ;

% Preallocating arrays
value_function = zeros( n_assets, max_age + 1 ) ;
labor_supplys  = zeros( n_assets, max_age     ) ;
savings        = zeros( n_assets, max_age     ) ;
consumptions   = zeros( n_assets, max_age     ) ;


% Solving the lifecycle model
for age = max_age:-1:1    % Instructing the for loop to iterate by increments of -1
    for ia = 1:n_assets
        continuation_values           = value_function( :, age+1 ) ;  % Continuation values over all possible savings
        asset                         = assets(ia)                 ;  % Solving at gridpoint for given asset
        initial_guess                 = [0.5, assets(ia) ];    % Making initial guess
        [maximizers, optimal_utility] = fminsearch(@(choice_space) calculate_utility(...
                                                       choice_space, asset , age, continuation_values,...
                                                       params ), initial_guess ,optim_options       );
        value_function( ia, age ) = -1*optimal_utility                                                                           ;
        labor_supplys ( ia, age ) = maximizers(1)                                                                                ;
        savings       ( ia, age ) = maximizers(2)                                                                                ;
        consumptions  ( ia, age ) = (1-params.tax_rate)*params.wage* maximizers(1) * params.productivity(age)...
                                      + params.interest_rate * asset  - maximizers(1)...
                                      + ge(age,params.ret_age)*params.benefit;
        
    end
    
end


savings_profile     = zeros( 1, max_age+1 ) ;    % Preallocating the savings profile
labor_profile       = zeros( 1, max_age   ) ;    % Preallocating the savings profile
consumption_profile = zeros( 1, max_age   ) ;
value_profile       = zeros( 1, max_age   ) ;

% Solving for optimal savings over the life-cycle, assuming that initial wealth is zero
savings_profile(1) = assets(1)            ;    % Assigning initial wealth to the first entry
for age = 1:max_age
    savings_profile     ( age+1 ) = interp1( assets, savings        ( :, age ), savings_profile(age), 'spline' );
    labor_profile       ( age   ) = interp1( assets, labor_supplys  ( :, age ), savings_profile(age), 'spline' );
    consumption_profile ( age   ) = interp1( assets, consumptions   ( :, age ), savings_profile(age), 'spline' );
    value_profile       ( age   ) = interp1( assets, value_function ( :, age ), savings_profile(age), 'spline' );
end
savings_profile = savings_profile(1:max_age);    % Removing superflous last period

% Creating aggregator function to keep the code DRY
aggregator = @(x) sum(x.*population_weights);

% Calculating and storing aggregates in structure
aggregates.consumption     =          aggregator( consumption_profile               );
aggregates.saving          =          aggregator( savings_profile                   );
aggregates.labor_supply    =          aggregator(                     labor_profile );
aggregates.labor_income    =          aggregator( wage*productivity'.*labor_profile );
aggregates.tax_revenue     = tax_rate*aggregator( wage*productivity'.*labor_profile );
aggregates.welfare         =          aggregator( value_profile                     );

aggregates.benefit         = benefit*sum( population_weights( params.ret_age:params.max_age ) );




function utility = calculate_utility( choice_space, asset, age, continuation_values, params )
    
    labor_supply = choice_space(1) ;
    saving       = choice_space(2) ;
    
    % Testing constraints
    if age<params.max_age
        minimum_saving = params.assets(1) ;    % Regular debt limit is lower bound of asset grid.
    elseif age==params.max_age
        minimum_saving = 0                ;    % Implementing no-Ponzi condition
    end
    
    % Treating saving constraint violation
    if labor_supply < 0 || labor_supply > 1
        utility = inf;
        return
    end
    
    if saving < minimum_saving || saving > params.assets(end)
        utility = inf;
        return
    end
            
    
    % Solving for consumption
    consumption = (1-params.tax_rate)*params.wage* labor_supply * params.productivity(age)...
                                      + params.interest_rate * asset  - saving...
                                      + ge(age,params.ret_age)*params.benefit;
    
    % Treating consumption constraint violation
    if consumption < 0
        utility = inf;
        return
    end
    
    % Solving for utility at the saving value
    continuation_value = interp1(params.assets, continuation_values, saving, 'spline' );
    
    utility = (1 /( 1-params.crra )) * ( (consumption^params.consumption_share) * (1-labor_supply)^(1-params.consumption_share) )^(1 - params.crra) + ...
               params.survival_probabilities(age) * params.discount_factor * continuation_value;
    utility = -1*utility ;
    
    
end

end