fire-planner/fire_planner/scenarios.py

"""Cartesian-product scenario generator.

Default counts:
  4 jurisdictions × 3 strategies × 5 leave-UK years × 2 glides = 120

Jurisdictions modelled by default: uk, nomad, cyprus, bulgaria.
Malaysia and Thailand are essentially equivalent in our tax engine
(both 0% on foreign income); pick one and document. Cyprus is
included because GeSY is non-trivial; Bulgaria for its 10% flat tax.

UK-stay scenarios duplicate across leave_uk_year (since you don't
leave) — kept in the product so the dashboard can present a uniform
heatmap; the simulator effectively ignores leave_year for UK.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from decimal import Decimal
from typing import Any

from fire_planner.glide_path import GLIDE_PATHS
from fire_planner.simulator import RegimeFn, constant_regime, jurisdiction_schedule
from fire_planner.strategies.base import WithdrawalStrategy
from fire_planner.strategies.guyton_klinger import GuytonKlingerStrategy
from fire_planner.strategies.spending_plan import SpendingPlanStrategy
from fire_planner.strategies.trinity import TrinityStrategy
from fire_planner.strategies.vpw import VpwStrategy, VpwWithFloorStrategy
from fire_planner.tax.base import TaxRegime
from fire_planner.tax.bulgaria import BulgariaTaxRegime
from fire_planner.tax.cyprus import CyprusTaxRegime
from fire_planner.tax.malaysia import MalaysiaTaxRegime
from fire_planner.tax.nomad import NomadTaxRegime
from fire_planner.tax.thailand import ThailandTaxRegime
from fire_planner.tax.uae import UaeTaxRegime
from fire_planner.tax.uk import UkTaxRegime

DEFAULT_JURISDICTIONS = ("uk", "nomad", "cyprus", "bulgaria")
DEFAULT_STRATEGIES = ("trinity", "guyton_klinger", "vpw")
DEFAULT_LEAVE_YEARS = (1, 2, 3, 4, 5)
DEFAULT_GLIDES = ("rising", "static_60_40")


@dataclass(frozen=True)
class ScenarioSpec:
    """One scenario in the Cartesian product."""
    jurisdiction: str
    strategy: str
    leave_uk_year: int
    glide_path: str
    spending_gbp: Decimal
    nw_seed_gbp: Decimal
    horizon_years: int = 60
    savings_per_year_gbp: Decimal = Decimal("0")
    config: dict[str, Any] = field(default_factory=dict)

    @property
    def external_id(self) -> str:
        return (f"{self.jurisdiction}-{self.strategy}-leave-y{self.leave_uk_year}-"
                f"glide-{self.glide_path}")


def build_strategy(
    name: str,
    floor: float | None = None,
    annual_real_adjust_pct: float = 0.0,
    guardrail_threshold_pct: float | None = None,
    guardrail_cut_pct: float = 0.10,
) -> WithdrawalStrategy:
    if name == "trinity":
        return TrinityStrategy()
    if name == "guyton_klinger":
        return GuytonKlingerStrategy()
    if name == "vpw":
        return VpwStrategy()
    if name == "vpw_floor":
        if floor is None:
            raise ValueError("vpw_floor strategy requires a `floor` value (real GBP)")
        return VpwWithFloorStrategy(floor=floor)
    if name == "custom":
        return SpendingPlanStrategy(
            annual_real_adjust_pct=annual_real_adjust_pct,
            guardrail_threshold_pct=guardrail_threshold_pct,
            guardrail_cut_pct=guardrail_cut_pct,
        )
    raise KeyError(f"Unknown strategy: {name!r}")


_JURISDICTION_CONSTRUCTORS: dict[str, type[TaxRegime]] = {
    "uk": UkTaxRegime,
    "nomad": NomadTaxRegime,
    "malaysia": MalaysiaTaxRegime,
    "thailand": ThailandTaxRegime,
    "cyprus": CyprusTaxRegime,
    "bulgaria": BulgariaTaxRegime,
    "uae": UaeTaxRegime,
}


def build_regime_schedule(jurisdiction: str, leave_uk_year: int) -> RegimeFn:
    """For UK-stay, returns a constant UK regime ignoring leave_year.
    For other jurisdictions, UK pre-departure and the target after."""
    if jurisdiction == "uk":
        return constant_regime(UkTaxRegime())
    cls = _JURISDICTION_CONSTRUCTORS.get(jurisdiction)
    if cls is None:
        raise KeyError(f"Unknown jurisdiction: {jurisdiction!r}")
    return jurisdiction_schedule(
        pre_departure=UkTaxRegime(),
        post_departure=cls(),
        leave_year=leave_uk_year,
    )


def cartesian_scenarios(
    spending_gbp: Decimal,
    nw_seed_gbp: Decimal,
    savings_per_year_gbp: Decimal = Decimal("0"),
    horizon_years: int = 60,
    jurisdictions: tuple[str, ...] = DEFAULT_JURISDICTIONS,
    strategies: tuple[str, ...] = DEFAULT_STRATEGIES,
    leave_years: tuple[int, ...] = DEFAULT_LEAVE_YEARS,
    glides: tuple[str, ...] = DEFAULT_GLIDES,
) -> list[ScenarioSpec]:
    out: list[ScenarioSpec] = []
    for jur in jurisdictions:
        for strat in strategies:
            for leave_y in leave_years:
                for glide in glides:
                    if glide not in GLIDE_PATHS:
                        raise KeyError(f"Unknown glide path: {glide!r}")
                    out.append(
                        ScenarioSpec(
                            jurisdiction=jur,
                            strategy=strat,
                            leave_uk_year=leave_y,
                            glide_path=glide,
                            spending_gbp=spending_gbp,
                            nw_seed_gbp=nw_seed_gbp,
                            horizon_years=horizon_years,
                            savings_per_year_gbp=savings_per_year_gbp,
                        ))
    return out