fire-planner/tests/test_simulator_events.py

"""End-to-end: cashflow_adjustments make portfolios bigger or smaller."""
from __future__ import annotations

import numpy as np
import pytest

from fire_planner.glide_path import static
from fire_planner.life_events import EventInput, events_to_cashflow_array
from fire_planner.simulator import simulate
from fire_planner.strategies.trinity import TrinityStrategy
from fire_planner.tax.malaysia import MalaysiaTaxRegime
from tests.test_simulator import fixed_paths


def _baseline_kwargs() -> dict[str, object]:
    """0% real returns, 25y, Trinity 4%, no taxes (Malaysia) — predictable."""
    paths = fixed_paths(n_paths=1, n_years=25, stock_ret=0.0, bond_ret=0.0, cpi=0.0)
    return dict(
        paths=paths,
        initial_portfolio=1_000_000.0,
        spending_target=40_000.0,
        glide=static(0.6),
        strategy=TrinityStrategy(initial_rate=0.04),
        regime=MalaysiaTaxRegime(),
    )


def test_no_adjustments_matches_baseline() -> None:
    base = simulate(**_baseline_kwargs())  # type: ignore[arg-type]
    with_zero = simulate(**_baseline_kwargs(), cashflow_adjustments=np.zeros(25))  # type: ignore[arg-type]
    np.testing.assert_allclose(base.portfolio_real, with_zero.portfolio_real)


def test_extra_outflows_show_up_in_withdrawal_trace() -> None:
    """A £100k spending bump in years 5-10 should be visible on the
    withdrawal trace — not just silently drained from the portfolio."""
    kwargs = _baseline_kwargs()
    adj = np.zeros(25, dtype=np.float64)
    extras = np.zeros(25, dtype=np.float64)
    adj[5:11] = -100_000.0  # drains the portfolio
    extras[5:11] = 100_000.0  # surfaces on the chart

    base = simulate(**kwargs)  # type: ignore[arg-type]
    bumped = simulate(**kwargs, cashflow_adjustments=adj, extra_outflows=extras)  # type: ignore[arg-type]

    # Year 0–4 unchanged (no extra outflow)
    np.testing.assert_allclose(base.withdrawal_real[:, :5], bumped.withdrawal_real[:, :5])
    # Years 5–10 should be ~100k higher than baseline (clipped only when
    # the portfolio was already drained — checked by spot-test).
    assert (bumped.withdrawal_real[:, 5:11] > base.withdrawal_real[:, 5:11]).all()
    # Year 5 specifically: strategy w (~40k) + 100k extra ≈ 140k.
    assert bumped.withdrawal_real[0, 5] == pytest.approx(140_000.0, rel=0.05)


def test_one_time_inheritance_lifts_portfolio() -> None:
    kwargs = _baseline_kwargs()
    adj = events_to_cashflow_array(
        [EventInput(year_start=10, one_time_amount_gbp=250_000)],
        horizon_years=25,
    )
    base = simulate(**kwargs)  # type: ignore[arg-type]
    enhanced = simulate(**kwargs, cashflow_adjustments=adj)  # type: ignore[arg-type]
    # Year 11 onward should be exactly £250k higher under 0% returns +
    # constant Trinity withdrawal.
    delta = enhanced.portfolio_real[0, 11:] - base.portfolio_real[0, 11:]
    assert np.all(delta > 0)
    # Year 11 specifically: +£250k landed at end of year 10, withdrawn
    # nothing extra in y10. By y11 just propagated forward.
    assert enhanced.portfolio_real[0, 11] - base.portfolio_real[0, 11] == 250_000


def test_ongoing_expense_drains_portfolio() -> None:
    kwargs = _baseline_kwargs()
    adj = events_to_cashflow_array(
        [EventInput(year_start=0, year_end=5, delta_gbp_per_year=-20_000)],
        horizon_years=25,
    )
    base = simulate(**kwargs)  # type: ignore[arg-type]
    drained = simulate(**kwargs, cashflow_adjustments=adj)  # type: ignore[arg-type]
    # 6 years × £20k expense = £120k less by end of year 6, 0% growth.
    delta = base.portfolio_real[0, 6] - drained.portfolio_real[0, 6]
    assert delta == 120_000


def test_event_can_force_failure() -> None:
    """A massive expense early on can ruin an otherwise-successful run."""
    kwargs = _baseline_kwargs()
    adj = events_to_cashflow_array(
        [EventInput(year_start=2, one_time_amount_gbp=-1_500_000)],
        horizon_years=25,
    )
    base = simulate(**kwargs)  # type: ignore[arg-type]
    ruined = simulate(**kwargs, cashflow_adjustments=adj)  # type: ignore[arg-type]
    assert base.success_rate == 1.0
    assert ruined.success_rate == 0.0