%% Set parameters
fra = 67; % full retirement age (born 1964 or after)
primary_pia = 1; % primary insurance amount
spouse_pia = [0.25 0.75]; % percentages of primary pia
%% Vector of ages
ages = 62:119; % ages 60 up to 119 (inclusive)
%% Compute early retirement benefit factors
early_benefit_factors = arrayfun(@(ages) benefit_factor_early(ages, 67), ...
    ages);
%% Compute delayed retirement benefit factors
delayed_benefit_factors = arrayfun(@(ages) benefit_factor_delayed(ages, 67), ...
    ages);
%% Compute spousal reduction factor
spousal_factors = arrayfun(@(age) spousal_factor(age, 67), ages);
%% Compute early and delayed retirement benefit amounts as percent of PIA
early_benefits = primary_pia .* (1 - early_benefit_factors) .* 100;
delayed_benefits = primary_pia .* (1 + delayed_benefit_factors) .* 100;
%% Concatenate early and delayed benefits
benefits = [early_benefits(1:6), delayed_benefits(7:9)];
%% Compute spousal benefits as percent of primary worker pia
spousal_benefit = (primary_pia/2) .* (1 - spousal_factors) .* 100;
%% Compute spouse's own benefit
spouse_early_benefits1 = spouse_pia(1) .* (1 - early_benefit_factors) .* 100;
spouse_delayed_benefits1 = spouse_pia(1) .* (1 + delayed_benefit_factors) .* 100;
spouse_own_benefit1 = [spouse_early_benefits1(1:6), spouse_delayed_benefits1(7:9)];
%% Compute spouse benefit (max of own benefit and spousal benefit)
spouse_benefit1 = max([spouse_own_benefit1; spousal_benefit(1:9)]);
%% Compute spouse's own benefit (different pia)
spouse_early_benefits2 = spouse_pia(2) .* (1 - early_benefit_factors) .* 100;
spouse_delayed_benefits2 = spouse_pia(2) .* (1 + delayed_benefit_factors) .* 100;
spouse_own_benefit2 = [spouse_early_benefits2(1:6), spouse_delayed_benefits2(7:9)];
%% Compute spouse benefit (different pia)
spouse_benefit2 = max([spouse_own_benefit2; spousal_benefit(1:9)]);
%% Plots
figure(1);
plot(ages(1:6), early_benefit_factors(1:6) .* 100, '-o');
title('Early Retirement Benefit Reduction Factor By Age (FRA = 67)');
xlabel('Claiming Age (Years)');
ylabel('Reduction Factor (Percent)');
saveas(figure(1), 'fig1.png');

figure(2)
plot(ages(6:9), delayed_benefit_factors(6:9) .* 100, '-o');
title('Delayed Retirement Benefit Increase Factor By Age (FRA = 67)');
xlabel('Claiming Age (Years)');
ylabel('Increase Factor (Percent)');
saveas(figure(2), 'fig2.png');

figure(3)
plot(ages(1:9), benefits, '-o')
title('Benefit Amounts Per Claiming Age as Percentages of PIA');
xlabel('Claiming Age (Years)');
ylabel('Percentage of PIA Received as Monthly Benefit');
saveas(figure(3), 'fig3.png');

figure(4);
plot(ages(1:6), spousal_factors(1:6) .* 100, '-o')
title('Spousal Benefit Reduction Factor By Claiming Age (FRA = 67)');
xlabel('Claiming Age (Years)');
ylabel('Reduction Factor (Percent)');
saveas(figure(4), 'fig4.png');

figure(5);
plot(ages(1:6), spousal_benefit(1:6), '-o');
title('Spousal Benefit By Claiming Age (FRA = 67)');
xlabel('Claiming Age (Years)');
ylabel("Spousal Benefit Received as Percentage of Primary Worker's PIA");
saveas(figure(5), 'fig5.png');

figure(6);
plot(ages(1:9), spouse_benefit1, '-o', ages(1:9), spouse_benefit2, '-o')
title('Spousal Benefit By Claiming Age For Different Spousal PIAs (FRA = 67)');
xlabel('Claiming Age (Years)');
ylabel("Spousal Benefit Received as Percentage of Primary Worker's PIA");
lgd = legend('25% of Primary PIA', '75% of Primary PIA');
title(lgd, "Spouse's PIA")
saveas(figure(6), 'fig6.png');

% Create mesh grid and totals for two different spousal PIAs
n = length(ages(1:9));
[Ah, As] = meshgrid(ages(1:9), ages(1:9));
total1 = zeros(n, n);
for i = 1:n
    for j = 1:n
        total1(i,j) = benefits(i) + spouse_benefit1(j);
    end
end

total2 = zeros(n, n);
for i = 1:n
    for j = 1:n
        total2(i,j) = benefits(i) + spouse_benefit2(j);
    end
end

figure(7);
surf(Ah, As, total1);
colorbar;
title("Benefits for Married Couples, Spouse's PIA is 25% of the Primary PIA (FRA = 67)");
ylabel('Primary Worker Claiming Age (Years)');
xlabel('Spouse Claiming Age (Years)');
zlabel("Monthly Benefits Received (Perentage of Primary Worker's PIA)")
saveas(figure(7), 'fig7.png');

figure(8);
surf(Ah, As, total2);
title("Benefits for Married Couples, Spouse's PIA is 75% of the Primary PIA (FRA = 67)");
colorbar;
ylabel('Primary Worker Claiming Age (Years)');
xlabel('Spouse Claiming Age (Years)');
zlabel("Monthly Benefits Received (Perentage of Primary Worker's PIA)")
saveas(figure(8), 'fig8.png');