fire-planner/tests/test_strategies.py

"""Withdrawal-strategy + glide-path behaviour."""
from fire_planner import glide_path
from fire_planner.strategies.base import StrategyState
from fire_planner.strategies.guyton_klinger import GuytonKlingerStrategy
from fire_planner.strategies.trinity import TrinityStrategy
from fire_planner.strategies.vpw import VpwStrategy, VpwWithFloorStrategy, pmt_rate


def state(**overrides: float | int) -> StrategyState:
    base = dict(
        portfolio=1_000_000.0,
        initial_portfolio=1_000_000.0,
        initial_withdrawal=40_000.0,
        year_idx=0,
        horizon_years=60,
        last_withdrawal=40_000.0,
        expected_real_return=0.04,
    )
    base.update(overrides)
    return StrategyState(**base)  # type: ignore[arg-type]


def test_trinity_year_zero_uses_initial_rate() -> None:
    s = TrinityStrategy(initial_rate=0.04)
    assert s.propose_withdrawal(state()) == 40_000.0


def test_trinity_holds_constant_in_real_terms() -> None:
    s = TrinityStrategy()
    assert s.propose_withdrawal(state(year_idx=10, last_withdrawal=40_000.0)) == 40_000.0


def test_trinity_doesnt_increase_with_portfolio_growth() -> None:
    s = TrinityStrategy()
    assert s.propose_withdrawal(state(year_idx=5, portfolio=2_000_000.0,
                                      last_withdrawal=40_000.0)) == 40_000.0


def test_gk_year_zero_uses_initial_withdrawal() -> None:
    """Year-0 honours the user's target spending (state.initial_withdrawal),
    not the strategy's preset rate. The preset rate now only matters as a
    fallback when initial_withdrawal isn't set."""
    s = GuytonKlingerStrategy(initial_rate=0.055)
    # state default has initial_withdrawal=40_000 → year 0 returns 40_000.
    assert s.propose_withdrawal(state()) == 40_000.0


def test_gk_year_zero_falls_back_to_preset_when_no_target() -> None:
    s = GuytonKlingerStrategy(initial_rate=0.055)
    # Override initial_withdrawal=0 → fall back to 5.5% × 1M = 55_000.
    assert s.propose_withdrawal(state(initial_withdrawal=0)) == 55_000.0


def test_gk_capital_preservation_cut() -> None:
    """Portfolio crashed: current rate now > 120% of the implied initial rate (5.5%);
    > 15y left → cut 10%. Implied rate = initial_withdrawal / initial_portfolio."""
    s = GuytonKlingerStrategy(initial_rate=0.055)
    # initial_withdrawal=55k, initial_portfolio=1M → implied rate = 5.5%.
    # last_w = 55k; portfolio = 700k → current rate = 7.86% > 6.6% guardrail.
    out = s.propose_withdrawal(state(year_idx=5, portfolio=700_000.0,
                                     last_withdrawal=55_000.0,
                                     initial_withdrawal=55_000.0))
    assert abs(out - 49_500.0) < 0.01


def test_gk_no_cut_when_horizon_under_15y_left() -> None:
    """Same crash, only 10y left — no cut applies."""
    s = GuytonKlingerStrategy(initial_rate=0.055)
    out = s.propose_withdrawal(
        state(year_idx=50, portfolio=700_000.0, last_withdrawal=55_000.0, horizon_years=60))
    assert out == 55_000.0


def test_gk_prosperity_bump() -> None:
    """Big bull market: current rate < 80% of 5.5% = 4.4% → bump 10%."""
    s = GuytonKlingerStrategy(initial_rate=0.055)
    out = s.propose_withdrawal(state(year_idx=5, portfolio=2_000_000.0, last_withdrawal=55_000.0))
    assert abs(out - 60_500.0) < 0.01


def test_pmt_rate_uniform_amortisation_at_zero_rate() -> None:
    assert abs(pmt_rate(years_remaining=60, real_rate=0.0) - 1 / 60) < 1e-12


def test_pmt_rate_full_drain_when_years_zero() -> None:
    assert pmt_rate(years_remaining=0, real_rate=0.04) == 1.0


def test_pmt_rate_bogleheads_table_60y() -> None:
    """Bogleheads VPW table: at 5% real, 60y, the published rate is
    5.28% (within £1/£10k of 5.2828% on a 60-year amortisation)."""
    assert abs(pmt_rate(60, 0.05) - 0.052828) < 1e-4


def test_pmt_rate_bogleheads_table_30y() -> None:
    """At 5% real, 30y → 6.51%."""
    assert abs(pmt_rate(30, 0.05) - 0.06505) < 1e-4


def test_pmt_rate_bogleheads_table_15y() -> None:
    """At 5% real, 15y → 9.63%."""
    assert abs(pmt_rate(15, 0.05) - 0.09634) < 1e-4


def test_vpw_year_zero_at_60y_horizon() -> None:
    """1M portfolio × pmt_rate(60, 0.05) = 1M × 0.0528 = 52,828.20."""
    s = VpwStrategy(expected_real_return=0.05)
    out = s.propose_withdrawal(state(horizon_years=60, year_idx=0))
    assert abs(out - 52_828.0) < 5  # within a few quid


def test_vpw_drain_at_horizon_end() -> None:
    """Last year: withdraw the entire portfolio."""
    s = VpwStrategy()
    out = s.propose_withdrawal(state(year_idx=59, horizon_years=60, portfolio=100_000.0))
    assert abs(out - 100_000.0) < 1


def test_vpw_with_floor_lifts_to_floor_when_vpw_proposes_less() -> None:
    """VPW on a 500k portfolio with 60y left at 5% would propose
    500k × 0.0528 ≈ 26,400. Floor=40k overrides — withdraw the floor."""
    s = VpwWithFloorStrategy(floor=40_000.0, expected_real_return=0.05)
    out = s.propose_withdrawal(state(portfolio=500_000.0, horizon_years=60, year_idx=0))
    assert out == 40_000.0


def test_vpw_with_floor_uses_vpw_when_above_floor() -> None:
    """VPW on a 2M portfolio with 60y left ≈ 105,656. Above floor=40k → use VPW."""
    s = VpwWithFloorStrategy(floor=40_000.0, expected_real_return=0.05)
    out = s.propose_withdrawal(state(portfolio=2_000_000.0, horizon_years=60, year_idx=0))
    assert abs(out - 105_656.0) < 50


def test_vpw_with_floor_clips_to_portfolio_when_portfolio_below_floor() -> None:
    """Terminal sequence: portfolio crashed below the floor — withdraw what's left."""
    s = VpwWithFloorStrategy(floor=40_000.0)
    out = s.propose_withdrawal(state(portfolio=15_000.0, horizon_years=60, year_idx=30))
    assert out == 15_000.0


def test_vpw_with_floor_zero_portfolio() -> None:
    s = VpwWithFloorStrategy(floor=40_000.0)
    out = s.propose_withdrawal(state(portfolio=0.0))
    assert out == 0.0


def test_vpw_with_floor_name() -> None:
    assert VpwWithFloorStrategy(floor=40_000.0).name == "vpw_floor"


def test_glide_rising_default_shape() -> None:
    g = glide_path.rising_equity()
    assert g(0) == 0.30
    assert abs(g(15) - 0.70) < 1e-9
    assert abs(g(30) - 0.70) < 1e-9
    # Halfway through the ramp
    assert abs(g(7) - (0.30 + 0.40 * 7 / 15)) < 1e-9


def test_glide_static() -> None:
    g = glide_path.static(0.60)
    assert g(0) == 0.60
    assert g(50) == 0.60


def test_glide_lookup() -> None:
    assert glide_path.get("rising")(0) == 0.30
    assert glide_path.get("static_60_40")(50) == 0.60