fire-planner/fire_planner/simulator.py

"""Monte Carlo simulator — the core of fire-planner.

Inputs:
- a `(n_paths, n_years, 3)` bootstrap of returns + CPI (`returns/`)
- a withdrawal strategy (`strategies/`)
- a glide-path function (`glide_path`)
- a tax regime (`tax/`)
- starting portfolio + spending target

Per path the simulator runs a 60-year (or whatever horizon) lifecycle:

  for each year y in 0..H:
      asset_alloc = glide(y)
      portfolio = portfolio * (1 + alloc·stock + (1-alloc)·bond)
      withdrawal = strategy.propose(state)
      tax = regime.compute_tax(...)
      portfolio -= (withdrawal + tax_in_addition)
      if portfolio < 0: failed_path

We work entirely in REAL GBP — convert returns to real (return/inflation
factor each year). Tax brackets are assumed to inflate with CPI (a v2
follow-up on fiscal drag).

Annual savings (during the accumulation phase) get added at year start
and earn the year's return.
"""
from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from decimal import Decimal

import numpy as np
import numpy.typing as npt

from fire_planner.flex_spending import FlexRule, applicable_cut
from fire_planner.glide_path import GlideFn
from fire_planner.strategies.base import StrategyState, WithdrawalStrategy
from fire_planner.tax.base import TaxInputs, TaxRegime

# Stock idx, bond idx, cpi idx in the bootstrap output.
STOCK = 0
BOND = 1
CPI = 2


@dataclass(frozen=True)
class SimulationResult:
    """Per-path arrays + scalar summaries.

    All money in REAL GBP (today's pounds). `portfolio_real[:, 0]` is
    the seed; `portfolio_real[:, k]` is end-of-year-k.
    """
    portfolio_real: npt.NDArray[np.float64]  # (n_paths, n_years+1)
    withdrawal_real: npt.NDArray[np.float64]  # (n_paths, n_years)
    tax_real: npt.NDArray[np.float64]  # (n_paths, n_years)
    success_mask: npt.NDArray[np.bool_]  # (n_paths,)

    @property
    def n_paths(self) -> int:
        return int(self.portfolio_real.shape[0])

    @property
    def n_years(self) -> int:
        return int(self.withdrawal_real.shape[1])

    @property
    def success_rate(self) -> float:
        return float(np.mean(self.success_mask))

    def ending_percentile(self, p: int) -> float:
        return float(np.percentile(self.portfolio_real[:, -1], p))

    def median_lifetime_tax(self) -> float:
        return float(np.median(self.tax_real.sum(axis=1)))

    def median_years_to_ruin(self) -> float | None:
        """Among failing paths, the median year-to-ruin (1-indexed).
        Returns None if every path survives."""
        failing = ~self.success_mask
        if not failing.any():
            return None
        portfolios = self.portfolio_real[failing, 1:]
        # First year-end where portfolio == 0 (or below)
        ruin_year = np.argmax(portfolios <= 0, axis=1) + 1
        return float(np.median(ruin_year))

    def sequence_risk_correlation(self) -> float:
        """Pearson correlation between year-1 drawdown and total
        success (1 if survived, 0 if not). Year-1 drawdown =
        (initial - portfolio_after_year1) / initial.

        Returns 0.0 when either variable has zero variance — e.g. all
        paths share the same year-1 returns (fixed_paths fixture) or
        every path succeeds.
        """
        initial = self.portfolio_real[:, 0]
        after_y1 = self.portfolio_real[:, 1]
        drawdown = (initial - after_y1) / initial
        success = self.success_mask.astype(np.float64)
        if np.var(drawdown) < 1e-12 or np.var(success) < 1e-12:
            return 0.0
        with np.errstate(invalid="ignore"):
            corr = np.corrcoef(drawdown, success)[0, 1]
        return 0.0 if np.isnan(corr) else float(corr)

    def fan_quantiles(self, p: int) -> npt.NDArray[np.float64]:
        """Per-year cross-path percentile of portfolio_real."""
        out: npt.NDArray[np.float64] = np.percentile(self.portfolio_real, p, axis=0)
        return out


# Default split of withdrawal across tax-treated buckets. The simulator
# treats withdrawals as "capital gains" by default since most accounts
# we model are taxable brokerage; ISA wraps don't tax at all and SIPP
# withdrawals are 25%-tax-free + ordinary income.
_BucketSplit = Callable[[float, int], TaxInputs]


def default_bucket_split(real_withdrawal: float, year_idx: int) -> TaxInputs:
    """Treat the entire withdrawal as long-term capital gains.
    Override via `bucket_split` to reflect ISA / SIPP / divs balances."""
    del year_idx
    return TaxInputs(capital_gains=Decimal(str(round(real_withdrawal, 2))))


RegimeFn = Callable[[int], TaxRegime]


def constant_regime(regime: TaxRegime) -> RegimeFn:
    return lambda _y: regime


def jurisdiction_schedule(
    pre_departure: TaxRegime,
    post_departure: TaxRegime,
    leave_year: int,
) -> RegimeFn:
    """While `year < leave_year` apply `pre_departure`; from `leave_year`
    onwards apply `post_departure`. Used to model "live in UK for N more
    years then move to Cyprus/Bulgaria/etc."
    """

    def fn(year: int) -> TaxRegime:
        return pre_departure if year < leave_year else post_departure

    return fn


def build_fixed_paths(
    n_paths: int,
    n_years: int,
    inflation_pct: float,
    stocks_growth_pct: float,
    stocks_dividend_pct: float,
    bonds_growth_pct: float,
    bonds_dividend_pct: float,
) -> npt.NDArray[np.float64]:
    """Build a deterministic ``(n_paths, n_years, 3)`` paths cube with
    constant nominal stock + bond returns and inflation. The nominal
    return per asset is ``growth + dividend``; the simulator's existing
    ``(1 + nominal) / (1 + cpi) − 1`` arithmetic then turns it into the
    real return path.

    Used by the Wave 1.D.3 fixed Rates mode — same shape as the Shiller
    bootstrap, so the rest of the engine doesn't change.
    """
    nominal_stock = stocks_growth_pct + stocks_dividend_pct
    nominal_bond = bonds_growth_pct + bonds_dividend_pct
    paths = np.zeros((n_paths, n_years, 3), dtype=np.float64)
    paths[:, :, STOCK] = nominal_stock
    paths[:, :, BOND] = nominal_bond
    paths[:, :, CPI] = inflation_pct
    return paths


def simulate(
    paths: npt.NDArray[np.float64],
    initial_portfolio: float,
    spending_target: float,
    glide: GlideFn,
    strategy: WithdrawalStrategy,
    regime: TaxRegime | RegimeFn,
    horizon_years: int | None = None,
    annual_savings: npt.NDArray[np.float64] | None = None,
    cashflow_adjustments: npt.NDArray[np.float64] | None = None,
    bucket_split: _BucketSplit = default_bucket_split,
    income_inflows: npt.NDArray[np.float64] | None = None,
    income_taxable: npt.NDArray[np.float64] | None = None,
    discretionary_outflows: npt.NDArray[np.float64] | None = None,
    flex_rules: list[FlexRule] | None = None,
) -> SimulationResult:
    """Run the MC simulation. `paths` shape: (n_paths, n_years, 3).

    `spending_target` is the year-0 real GBP draw; subsequent years are
    decided by the strategy. `annual_savings`, if given, is a (n_years,)
    real-GBP array — added at the start of each year while accumulating.

    `cashflow_adjustments`, if given, is a (n_years,) real-GBP array of
    per-year deltas applied **after** savings + return but **before**
    withdrawal. Positive = inflow (e.g. inheritance, rental income),
    negative = extra outflow (e.g. childcare, sabbatical). Used to plumb
    `life_event` rows into the projection.

    `regime` may be a single `TaxRegime` (constant for all years) or a
    callable `(year_idx) -> TaxRegime` to model jurisdiction switches —
    e.g. UK for years 0..N-1, then Cyprus from year N onward.
    """
    regime_at: RegimeFn = (regime if callable(regime) else constant_regime(regime))
    n_paths, n_years, _ = paths.shape
    if horizon_years is None:
        horizon_years = n_years

    portfolio = np.full(n_paths, float(initial_portfolio), dtype=np.float64)
    portfolio_history = np.zeros((n_paths, n_years + 1), dtype=np.float64)
    portfolio_history[:, 0] = portfolio
    withdrawal_hist = np.zeros((n_paths, n_years), dtype=np.float64)
    tax_hist = np.zeros((n_paths, n_years), dtype=np.float64)
    last_withdrawal = np.full(n_paths, float(spending_target), dtype=np.float64)

    if annual_savings is None:
        annual_savings = np.zeros(n_years, dtype=np.float64)
    if cashflow_adjustments is None:
        cashflow_adjustments = np.zeros(n_years, dtype=np.float64)
    if income_inflows is None:
        income_inflows = np.zeros(n_years, dtype=np.float64)
    if income_taxable is None:
        income_taxable = np.zeros(n_years, dtype=np.float64)
    if discretionary_outflows is None:
        discretionary_outflows = np.zeros(n_years, dtype=np.float64)
    rules = list(flex_rules) if flex_rules else []
    # Track running ATH per path so we can decide flex cuts each year.
    ath = np.full(n_paths, float(initial_portfolio), dtype=np.float64)

    for y in range(n_years):
        alloc = glide(y)
        # Real returns for this year, all paths: shape (n_paths,)
        nominal_stock = paths[:, y, STOCK]
        nominal_bond = paths[:, y, BOND]
        cpi = paths[:, y, CPI]
        real_stock = (1 + nominal_stock) / (1 + cpi) - 1
        real_bond = (1 + nominal_bond) / (1 + cpi) - 1
        port_return = alloc * real_stock + (1 - alloc) * real_bond

        # Add savings + recurring income inflows at year start, apply
        # year's return, then apply life-event cashflow adjustments.
        # Adjustments don't compound this year's returns (they're
        # treated as end-of-year events). Income tax on `income_taxable[y]`
        # is collected once outside the per-path loop below — it's path
        # invariant when the regime is the same for all paths.
        portfolio = (portfolio + annual_savings[y] + income_inflows[y]) * (1 + port_return)
        portfolio = portfolio + cashflow_adjustments[y]
        if income_taxable[y] > 0.0:
            income_tax_breakdown = regime_at(y).compute_tax(
                TaxInputs(earned_income=Decimal(str(round(float(income_taxable[y]), 2)))))
            portfolio = portfolio - float(income_tax_breakdown.total)
        # Flex spending: per-path, decide the cut from this year's
        # drawdown-from-ATH and refund the trimmed discretionary
        # back to the portfolio. The cashflow_adjustments array already
        # subtracted the *baseline* discretionary, so we add back
        # `cut_pct * baseline` to leave only the post-cut amount drawn.
        if rules and discretionary_outflows[y] > 0.0:
            for p in range(n_paths):
                if ath[p] <= 0:
                    continue
                drawdown = max(0.0, 1.0 - portfolio[p] / ath[p])
                cut = applicable_cut(drawdown, rules)
                if cut > 0:
                    portfolio[p] += cut * float(discretionary_outflows[y])

        # Strategy is per-path Python — 600k iterations at 60y × 10k paths.
        # Profiled: ~3 seconds for the full Trinity / GK / VPW set.
        for p in range(n_paths):
            state = StrategyState(
                portfolio=float(portfolio[p]),
                initial_portfolio=float(initial_portfolio),
                initial_withdrawal=float(spending_target),
                year_idx=y,
                horizon_years=horizon_years,
                last_withdrawal=float(last_withdrawal[p]),
            )
            w = strategy.propose_withdrawal(state)
            # Strategies are real-GBP; clip to portfolio.
            w = max(0.0, min(w, float(portfolio[p])))
            tax_breakdown = regime_at(y).compute_tax(bucket_split(w, y))
            t = float(tax_breakdown.total)
            # Drain BOTH withdrawal AND tax from the portfolio. Mental
            # model: `w` is what the user takes home to spend; the tax
            # is an additional drag that comes out of the same pool. A
            # high-tax jurisdiction therefore drains the portfolio
            # faster and lowers the success rate, matching user intuition
            # (Cyprus non-dom should beat UK on outcome, not just on a
            # summary stat). Pre-2026-05-10 the engine recorded `t` but
            # didn't subtract it, so jurisdiction only changed the
            # `median_lifetime_tax_gbp` cell while the fan chart and
            # success rate were identical across regimes.
            portfolio[p] = max(0.0, portfolio[p] - w - t)
            withdrawal_hist[p, y] = w
            tax_hist[p, y] = t
            last_withdrawal[p] = w

        portfolio_history[:, y + 1] = np.clip(portfolio, a_min=0.0, a_max=None)
        portfolio = portfolio_history[:, y + 1]
        # Update running ATH per path so next year's flex decision uses
        # the post-close peak.
        np.maximum(ath, portfolio, out=ath)

    # Success = portfolio stayed positive through every interim year.
    # Excludes the very last year-end because VPW deliberately drains
    # to ~0 at the horizon boundary by construction; that's not a failure.
    success_mask = portfolio_history[:, 1:-1].min(axis=1) > 0.0
    return SimulationResult(
        portfolio_real=portfolio_history,
        withdrawal_real=withdrawal_hist,
        tax_real=tax_hist,
        success_mask=success_mask,
    )