%% Week 5 Challenge (Solution) — Labor Market Equilibrium (Exogenous Wage)
% Solve Z(h) = 0 for h in (0,1)
% Methods: (i) Bisection, (ii) Damped Newton
% Then: comparative statics over tau in [0, 0.4]
% This version uses a partial-equilibrium household FOC with exogenous wage.
% Parameters chosen to guarantee a root on (0,1).

clear; clc; close all;

%% --- PARAMETERS ---
w     = 1.0;     % exogenous wage
tau   = 0.20;    % tax rate
a     = 0.50;    % lump-sum transfer
g     = 0.00;    % government spending (ignored in c here)
sigma = 2.0;     % CRRA coefficient
phi   = 2.0;     % leisure curvature
chi   = 3.7;   % weight on disutility of work

% Domain for h
h_low  = 1e-6;
h_high = 1 - 1e-6;

%% --- MODEL OBJECTS ---
% Budget (short run): c(h) = (1 - tau) * w * h + a - g
c  = @(h) (1 - tau) * w .* h + a - g;

% Household FOC: U_c(c(h)) * (1 - tau) * w = V_l(1 - h)
% U_c(c(h)) * (1 - tau) * w - V_l(1 - h) = 0
% U_c(c) = c^(-sigma),  V_l(ell) = chi * ell^phi
Z  = @(h) c(h).^(-sigma) * (1 - tau) * w - chi * (1 - h).^phi;

% Quick check: endpoints (ensures sign change inside (0,1))
Zlow = Z(h_low);
Zhigh = Z(h_high);
fprintf('Z at endpoints: Z(%.1e)=%.3f, Z(1-%.1e)=%.3f\n', h_low, Zlow, 1-h_high, Zhigh);

%% --- PART 1: BISECTION (robust) ---
tol   = 1e-8;
maxit = 100;

% We expect Z(h_low) < 0 and Z(h_high) > 0 with the chosen parameters.
assert(Zlow * Zhigh < 0, 'No sign change on (0,1). Adjust parameters.');

low = h_low; high = h_high;
it_bis = 0;
while it_bis < maxit
    it_bis = it_bis + 1;
    mid  = 0.5 * (low + high);
    Zmid = Z(mid);
    if abs(Zmid) < tol || (high - low) < tol
        break
    end
    if Zlow * Zmid < 0
        high = mid;    % root in [low, mid]
        Zhigh = Zmid;
    else
        low  = mid;    % root in [mid, high]
        Zlow = Zmid;
    end
end
h_bis = mid;
fprintf('Bisection: h* = %.6f, |Z(h*)| = %.2e, it = %d\n', h_bis, abs(Z(h_bis)), it_bis);

%% --- PART 1: DAMPED NEWTON (finite-difference derivative) ---
fd = @(f,x) (f(x + 1e-6) - f(x - 1e-6)) / (2e-6);

h = 0.50;               % good starting guess inside (0,1)
alpha_step = 0.5;       % damping
it_new = 0;

for it = 1:maxit
    it_new = it;
    Zh = Z(h);
    if abs(Zh) < tol
        break
    end
    dZh = fd(Z, h);
    if ~isfinite(dZh) || abs(dZh) < 1e-12
        warning('Derivative problem at iteration %d (dZ ~ 0).', it);
        break
    end
    h = h - alpha_step * Zh / dZh;
    % keep in domain
    h = min(max(h, h_low), h_high);
end
h_new = h;
fprintf('Newton (damped): h* = %.6f, |Z(h*)| = %.2e, it = %d\n', h_new, abs(Z(h_new)), it_new);

%% --- PART 2: COMPARATIVE STATICS OVER tau ---
taus = linspace(0, 0.4, 9);
H_tau = nan(size(taus));

for j = 1:numel(taus)
    tau_j = taus(j);
    c_j   = @(h) (1 - tau_j) * w .* h + a - g;
    Z_j   = @(h) c_j(h).^(-sigma) * (1 - tau_j) * w - chi * (1 - h).^phi;

    % Recompute endpoints and ensure sign change; if not, mark NaN
    Zl = Z_j(h_low); Zh = Z_j(h_high);
    if Zl * Zh > 0
        % No interior root for this tau_j with these parameters
        H_tau(j) = NaN;
        continue
    end

    low = h_low; high = h_high;
    for it = 1:80
        mid = 0.5 * (low + high);
        Zm  = Z_j(mid);
        if abs(Zm) < 1e-8 || (high - low) < 1e-8
            break
        end
        if Z_j(low) * Zm < 0
            high = mid;
        else
            low = mid;
        end
    end
    H_tau(j) = mid;
end

% Plot h*(tau)
figure('Color','w');
plot(taus, H_tau, '-o', 'LineWidth', 1.2, 'MarkerSize', 6);
grid on; box on;
xlabel('\tau'); ylabel('h^*(\tau)');
title('Equilibrium Hours vs. Tax Rate');

mkdir('Figures')
% OPTIONAL: save figure
exportgraphics(gcf, 'Figures/week5_challenge_h_vs_tau.png', 'Resolution', 300);

%% --- Visual check of Z(h) at baseline ---
hs = linspace(h_low, h_high, 400);
figure('Color','w');
plot(hs, arrayfun(Z, hs), 'LineWidth', 1.2); grid on; box on;
yline(0, ':k'); xlabel('h'); ylabel('Z(h)');
title(sprintf('Baseline Z(h) with tau=%.2f', tau));
