%% Week 5 Homework (Solution)
% Solving nonlinear equations in macroeconomics
% - Exercise 1: IS–LM with nonlinear investment and money demand (system)
% - Exercise 2: Labor supply with exogenous wage (scalar root) — vary sigma
% - Exercise 3: Convergence comparison (Bisection vs Damped Newton)
%
% Author: (solution key)
% -------------------------------------------------------------------------

clear; clc; close all;

fprintf('\n==================== EXERCISE 1: Nonlinear IS–LM ====================\n');

%% EXERCISE 1 — Nonlinear IS–LM (system in Y and i>0 via reparam i = exp(s))
% Model:
%   Y = C0 + c*(Y - T) + I0 - beta*i^2 + G
%   M/P = k*Y - lambda*i
% with i = exp(s)  (guarantees i > 0)

% Parameters
C0    = 50;
c     = 0.8;
I0    = 60;
betaI = 2;        % curvature in investment: I(i) = I0 - beta*i^2
T     = 50;
G     = 100;
MP    = 400;      % M/P
k     = 0.5;
lambda= 20;

% System in [Y; s], where i = exp(s) > 0
Fpos = @(ys) [ ...
   ys(1) - (C0 + c*(ys(1) - T) + I0 - betaI*exp(2*ys(2)) + G);  % IS: i^2 = exp(2s)
   k*ys(1) - lambda*exp(ys(2)) - MP                             % LM:  i = exp(s)
];

% Two initial guesses (Y0, s0) with positive i = exp(s0)
x0a = [500; log(3)];   % i0 = 3  > 0
x0b = [300; log(8)];   % i0 = 8  > 0

% Use fsolve if available; otherwise, simple damped Newton on [Y;s]
use_fsolve = exist('fsolve','file') == 2;

if use_fsolve
    opts = optimoptions('fsolve','Display','off','FunctionTolerance',1e-12,...
                        'StepTolerance',1e-12,'MaxIterations',400);
    [ya, fa, exitflaga, outa] = fsolve(Fpos, x0a, opts);
    [yb, fb, exitflagb, outb] = fsolve(Fpos, x0b, opts);

    ia = exp(ya(2));   % recover i from s
    ib = exp(yb(2));

    fprintf('Start A: [Y0=%.1f, i0=%.1f] -> (Y*, i*) = (%.6f, %.6f), ||F||=%.2e, it=%d\n', ...
        x0a(1), exp(x0a(2)), ya(1), ia, norm(fa), outa.iterations);
    fprintf('Start B: [Y0=%.1f, i0=%.1f] -> (Y*, i*) = (%.6f, %.6f), ||F||=%.2e, it=%d\n', ...
        x0b(1), exp(x0b(2)), yb(1), ib, norm(fb), outb.iterations);
else
    fprintf('fsolve not found — using a simple damped Newton fallback (in [Y;s]).\n');
    ya = newton_system(@(z) Fpos(z), x0a, 1e-12, 200);
    yb = newton_system(@(z) Fpos(z), x0b, 1e-12, 200);
    ia = exp(ya(2)); ib = exp(yb(2));
    fprintf('Start A -> (Y*, i*) = (%.6f, %.6f), ||F||=%.2e\n', ya(1), ia, norm(Fpos(ya)));
    fprintf('Start B -> (Y*, i*) = (%.6f, %.6f), ||F||=%.2e\n', yb(1), ib, norm(Fpos(yb)));
end

% ---- Visualization: LM and ONLY the positive IS branch ----
Ygrid = linspace(100, 1000, 400);

% LM: kY - lambda i = M/P  =>  i_LM = (kY - MP)/lambda, but we clip at 0.
i_LM  = max(0, (k*Ygrid - MP)/lambda);

% IS (positive branch only): beta*i^2 = C0 + c(Y-T) + I0 + G - Y
RHS   = (C0 + c*(Ygrid - T) + I0 + G - Ygrid)/betaI;
i_IS  = sqrt(max(0, RHS));   % keep only i >= 0 branch

figure('Color','w');
plot(Ygrid, i_LM, 'LineWidth', 1.6, 'DisplayName','LM (i \geq 0)'); hold on; grid on; box on; 
plot(Ygrid, i_IS, 'LineWidth', 1.6, 'DisplayName','IS (positive branch)');
plot(ya(1), ia, 'ro', 'MarkerFaceColor','r', 'DisplayName','Equilibrium (start A)');
plot(yb(1), ib, 'ko', 'MarkerFaceColor','k', 'DisplayName','Equilibrium (start B)');
xlabel('Output Y'); ylabel('Interest rate i (>0)');
title('Nonlinear IS–LM Equilibrium (Positive-Rate Branch)');
legend('Location','best');
exportgraphics(gcf, 'Figures/IS_LM.png', 'Resolution', 300);

% -------------------------------------------------------------------------
fprintf('\n=================== EXERCISE 2: Labor Supply (scalar) ===================\n');

%% EXERCISE 2 — Labor supply for varying sigma (Bisection)
% Setup matches the lecture/challenge: exogenous wage.
% Z(h) = c(h)^(-sigma) * (1 - tau) * w - chi * (1 - h)^phi = 0,
% c(h) = (1 - tau) * w * h + a - g

w     = 1.0;
tau   = 0.20;
a     = 0.50;
g     = 0.00;
phi   = 2.0;
chi   = 3.7;

sigmas = 1:5;
Hsig   = nan(size(sigmas));

for s = 1:numel(sigmas)
    sigma = sigmas(s);
    cfun  = @(h) (1 - tau) * w .* h + a - g;
    Z     = @(h) cfun(h).^(-sigma) * (1 - tau) * w - chi * (1 - h).^phi;

    % Robust bisection on (0,1)
    low  = 1e-6; high = 1 - 1e-6;
    Zlow = Z(low); Zhigh = Z(high);
    if Zlow * Zhigh > 0
        % If no sign change, record NaN (corner solution h*=1 economically)
        % For sigma >= 3, Z(h) may not cross zero in (0,1): very high risk
        % aversion makes leisure dominate consumption in all interior points.
        % The FOC Z(h)=0 has no interior root, so h*=1 (work full time).
        Hsig(s) = NaN;
        continue
    end
    for it = 1:100
        mid = 0.5*(low + high);
        Zm  = Z(mid);
        if abs(Zm) < 1e-10 || (high - low) < 1e-10
            break
        end
        if Zlow * Zm < 0
            high = mid;
            Zhigh = Zm;
        else
            low  = mid;
            Zlow = Zm;
        end
    end
    Hsig(s) = mid;
    fprintf('sigma=%d -> h* = %.6f, |Z|=%.1e\n', sigma, mid, abs(Z(mid)));
end

% Plot h*(sigma)
% Note: For sigma >= 3, no interior solution exists (corner h*=1),
% so plot only shows valid points (NaN values are skipped)
figure('Color','w');
plot(sigmas, Hsig, '-o', 'LineWidth', 1.6, 'MarkerSize', 7);
grid on; box on;
xlabel('\sigma'); ylabel('h^*(\sigma)');
title('Labor Supply vs. Risk Aversion (\sigma)');
exportgraphics(gcf, 'Figures/LS_RA.png', 'Resolution', 300);

% Print summary
fprintf('\nSummary: For sigma=1,2: interior solution exists (h* < 1)\n');
fprintf('        For sigma>=3: corner solution (h*=1) - no interior root\n');

%A higher \sigma means the household's marginal utility of consumption increases more sharply as consumption falls — the agent becomes more "cautious" and less willing to trade leisure for uncertain or small consumption gains.
%	•	Higher \sigma: consumption is valued less at the margin, so leisure
%	becomes relatively more attractive.
%	•	Lower \sigma: flatter utility in c, stronger incentive to work for extra income.

%Hence, in this static model, risk aversion acts like a "leisure preference amplifier" — pushing equilibrium hours downward as \sigma increases.
% -------------------------------------------------------------------------
fprintf('\n============= EXERCISE 3: Convergence speed (scalar problem) ============\n');

%% EXERCISE 3 — Convergence comparison: Bisection vs Damped Newton
% Reuse baseline Z(h) with sigma=2
sigma = 2.0;
cfun  = @(h) (1 - tau) * w .* h + a - g;
Z     = @(h) cfun(h).^(-sigma) * (1 - tau) * w - chi * (1 - h).^phi;

% Target tolerance
tol = 1e-10;
maxit = 200;

% Bisection from full bracket
[hb, it_bis] = bisection_solve(Z, 1e-6, 1-1e-6, tol, maxit);

% Newton from various starting guesses
starts = [0.2, 0.5, 0.8];
alpha  = 0.5;

results = [];  % for printing table
for s = 1:numel(starts)
    [hn, it_new] = newton_damped_scalar(Z, starts(s), alpha, tol, maxit, 1e-6, 1-1e-6);
    results = [results; {"Newton (damped)", starts(s), it_new, hn, abs(Z(hn))}]; %#ok<AGROW>
end
results = [ {"Bisection", NaN, it_bis, hb, abs(Z(hb))}; results ];

% Print table
fprintf('\n%-16s | %14s | %10s | %12s | %12s\n', 'Method','Start guess','Iterations','h*','|Z(h*)|');
fprintf('%s\n', repmat('-',1,70));
for r = 1:size(results,1)
    meth = results{r,1};
    s0   = results{r,2};
    it   = results{r,3};
    hst  = results{r,4};
    res  = results{r,5};
    if isnan(s0)
        fprintf('%-16s | %14s | %10d | %12.6f | %12.1e\n', meth, '—', it, hst, res);
    else
        fprintf('%-16s | %14.2f | %10d | %12.6f | %12.1e\n', meth, s0, it, hst, res);
    end
end

% ============================== HELPERS ==============================
function [root, it] = bisection_solve(Z, low, high, tol, maxit)
% Robust bisection on [low, high]
Zlow  = Z(low);  Zhigh = Z(high);
if Zlow*Zhigh > 0
    error('Bisection: no sign change on the interval.');
end
for it = 1:maxit
    mid = 0.5*(low + high);
    Zm  = Z(mid);
    if abs(Zm) < tol || (high - low) < tol
        root = mid; return;
    end
    if Zlow * Zm < 0
        high = mid; Zhigh = Zm;
    else
        low  = mid; Zlow  = Zm;
    end
end
root = mid;
end

function [h, it] = newton_damped_scalar(Z, h0, alpha, tol, maxit, h_low, h_high)
% Damped Newton for scalar root with finite-diff derivative and projection
fd = @(f,x) (f(x+1e-6) - f(x-1e-6)) / (2e-6);
h = h0;
for it = 1:maxit
    Zh = Z(h);
    if abs(Zh) < tol, break; end
    dZh = fd(Z, h);
    if ~isfinite(dZh) || abs(dZh) < 1e-14
        % fall back to a tiny bisection step near current h
        h = min(max(h + 1e-3*sign(-Zh), h_low), h_high);
        continue
    end
    h = h - alpha * Zh / dZh;
    % project to domain
    h = min(max(h, h_low), h_high);
end
end