fire-planner/fire_planner/goals_eval.py

"""Evaluate retirement goals against Monte Carlo simulation results.

Goal-kind contract (v1):

- ``target_nw_by_year`` — at year ``target_year`` the real portfolio must
  satisfy ``comparator target_amount_gbp`` (e.g. ``>= 2_000_000``).
  Probability = fraction of paths that hit the comparator at that year.
- ``never_run_out`` — portfolio must stay > 0 at every year up to
  ``target_year`` (or the full horizon if ``target_year`` is None).
  Probability = fraction of paths that never hit zero in the window.
- ``target_real_income`` — median (over the window ``target_year`` ..
  horizon) real withdrawal must satisfy ``comparator target_amount_gbp``.
  Probability = fraction of paths whose median window withdrawal hits.

Goal kinds the v1 evaluator does not yet recognise return probability=0.0
and ``passed=False`` so the API surface stays uniform — frontend can
render an "unsupported kind" hint.
"""
from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass
from decimal import Decimal
from typing import Protocol

import numpy as np

from fire_planner.simulator import SimulationResult


class _GoalLike(Protocol):
    """Duck-typed goal — ORM `RetirementGoal` and unsaved goals both fit."""
    kind: str
    name: str
    target_amount_gbp: Decimal | None
    target_year: int | None
    comparator: str
    success_threshold: Decimal


@dataclass(frozen=True)
class GoalEvaluation:
    goal_id: int | None
    name: str
    kind: str
    probability: float
    threshold: float
    passed: bool


_COMPARATORS = {
    ">=": np.greater_equal,
    ">": np.greater,
    "<=": np.less_equal,
    "<": np.less,
    "=": np.equal,
}


def _comparator_fn(op: str):  # type: ignore[no-untyped-def]
    return _COMPARATORS.get(op, np.greater_equal)


def _eval_one(
    result: SimulationResult,
    goal: _GoalLike,
    horizon_years: int,
) -> tuple[float, bool]:
    threshold = float(goal.success_threshold)

    if goal.kind == "target_nw_by_year":
        if goal.target_year is None or goal.target_amount_gbp is None:
            return 0.0, False
        # portfolio_real has columns 0..n_years (year 0 = seed). Clip.
        col = max(0, min(int(goal.target_year), horizon_years))
        target = float(goal.target_amount_gbp)
        cmp_fn = _comparator_fn(goal.comparator)
        hits = cmp_fn(result.portfolio_real[:, col], target)
        prob = float(np.mean(hits))
        return prob, prob >= threshold

    if goal.kind == "never_run_out":
        # Stay > 0 across years 1..target_year (excludes seed year 0).
        end = int(goal.target_year) if goal.target_year is not None else horizon_years
        end = max(1, min(end, horizon_years))
        # portfolio_real has n_years+1 cols; index 1..end inclusive.
        survived = (result.portfolio_real[:, 1:end + 1] > 0.0).all(axis=1)
        prob = float(np.mean(survived))
        return prob, prob >= threshold

    if goal.kind == "target_real_income":
        if goal.target_amount_gbp is None:
            return 0.0, False
        start_y = int(goal.target_year) if goal.target_year is not None else 0
        start_y = max(0, min(start_y, horizon_years - 1))
        window = result.withdrawal_real[:, start_y:]
        if window.size == 0:
            return 0.0, False
        path_medians = np.median(window, axis=1)
        target = float(goal.target_amount_gbp)
        cmp_fn = _comparator_fn(goal.comparator)
        hits = cmp_fn(path_medians, target)
        prob = float(np.mean(hits))
        return prob, prob >= threshold

    return 0.0, False


def evaluate_goals(
    result: SimulationResult,
    goals: Iterable[_GoalLike],
    horizon_years: int | None = None,
) -> list[GoalEvaluation]:
    """Compute probability + pass/fail for every enabled goal."""
    H = horizon_years if horizon_years is not None else result.n_years
    out: list[GoalEvaluation] = []
    for goal in goals:
        if not getattr(goal, "enabled", True):
            continue
        prob, passed = _eval_one(result, goal, H)
        out.append(
            GoalEvaluation(
                goal_id=getattr(goal, "id", None),
                name=goal.name,
                kind=goal.kind,
                probability=prob,
                threshold=float(goal.success_threshold),
                passed=passed,
            ))
    return out