fire-planner/fire_planner/api/spending_profile.py

"""Per-year spending breakdown — drives the Visx stacked-area chart on
the Plan tab.

Returns one ``SpendingProfilePoint`` per year of the scenario horizon:

  base_gbp           — scenario-level baseline spending (real GBP)
  essential_gbp      — sum of |delta| from active essential life events
  discretionary_gbp  — sum of |delta| from active discretionary events
  not_spending_gbp   — informational events that have a delta_gbp_per_year
                       but are tagged ``not_spending`` (rare, but possible)
  flex_cut_gbp       — discretionary £ trimmed by flex rules at p50 portfolio
                       drawdown vs running ATH (0 when no rules / no
                       drawdown).
  total_gbp          — base + essential + discretionary − flex_cut

The flex-cut estimate uses the persisted p50 portfolio path (from
``projection_yearly``) because the user's scenario.config_json may carry
``flex_rules`` they've configured but not yet re-run a recompute against.
This gives an honest preview of how the rules would shape spending; the
exact per-path cut shows up in the next live ``POST /simulate``.
"""
from __future__ import annotations

from collections.abc import Sequence
from decimal import Decimal

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from fire_planner.api.dependencies import get_session
from fire_planner.api.schemas import (
    SpendingProfilePoint,
    SpendingProfilePreviewRequest,
    SpendingProfileResponse,
)
from fire_planner.db import LifeEvent, McRun, ProjectionYearly, Scenario

router = APIRouter(prefix="/scenarios", tags=["spending-profile"])


def _category_outflow_at(events: Sequence[LifeEvent], year_idx: int,
                         category: str) -> Decimal:
    total = Decimal("0")
    for ev in events:
        if not ev.enabled or ev.category != category:
            continue
        if year_idx < ev.year_start:
            continue
        end = ev.year_end if ev.year_end is not None else ev.year_start
        if year_idx > end:
            continue
        delta = Decimal(str(ev.delta_gbp_per_year or 0))
        if delta < 0:
            total += -delta
        if year_idx == ev.year_start and ev.one_time_amount_gbp is not None:
            ot = Decimal(str(ev.one_time_amount_gbp))
            if ot < 0:
                total += -ot
    return total


def _category_inflow_at(events: Sequence[LifeEvent], year_idx: int,
                        category: str) -> Decimal:
    """Positive deltas are inflows — surface them as a negative spending
    contribution so the stacked area sums correctly."""
    total = Decimal("0")
    for ev in events:
        if not ev.enabled or ev.category != category:
            continue
        if year_idx < ev.year_start:
            continue
        end = ev.year_end if ev.year_end is not None else ev.year_start
        if year_idx > end:
            continue
        delta = Decimal(str(ev.delta_gbp_per_year or 0))
        if delta > 0:
            total += delta
    return total


def _flex_rules_for(scenario: Scenario) -> list[dict[str, float]]:
    blob = scenario.config_json or {}
    rules = blob.get("flex_rules") if isinstance(blob, dict) else None
    if not isinstance(rules, list):
        return []
    out: list[dict[str, float]] = []
    for r in rules:
        if not isinstance(r, dict):
            continue
        try:
            out.append({
                "from_ath_pct": float(r.get("from_ath_pct", 0)),
                "cut_discretionary_pct": float(r.get("cut_discretionary_pct", 0)),
            })
        except (TypeError, ValueError):
            continue
    return out


def _flex_cut_at_year(year: ProjectionYearly, ath_so_far_gbp: Decimal,
                      rules: list[dict[str, float]],
                      discretionary_gbp: Decimal) -> Decimal:
    if not rules or discretionary_gbp <= 0 or ath_so_far_gbp <= 0:
        return Decimal("0")
    p50 = Decimal(str(year.p50_portfolio_gbp))
    drawdown = max(Decimal("0"), Decimal("1") - p50 / ath_so_far_gbp)
    best = Decimal("0")
    for rule in rules:
        thr = Decimal(str(rule["from_ath_pct"]))
        cut = Decimal(str(rule["cut_discretionary_pct"]))
        if drawdown >= thr and cut > best:
            best = cut
    return (discretionary_gbp * best).quantize(Decimal("0.01"))


@router.get("/{scenario_id}/spending-profile",
            response_model=SpendingProfileResponse)
async def get_spending_profile(
    scenario_id: int,
    session: AsyncSession = Depends(get_session),
) -> SpendingProfileResponse:
    scen = await session.get(Scenario, scenario_id)
    if scen is None:
        raise HTTPException(status_code=404, detail="Scenario not found")

    events = list((await session.execute(
        select(LifeEvent).where(LifeEvent.scenario_id == scenario_id))).scalars().all())

    yearly_rows: list[ProjectionYearly] = []
    run = (await session.execute(
        select(McRun).where(McRun.scenario_id == scenario_id).order_by(
            McRun.run_at.desc()).limit(1))).scalar_one_or_none()
    if run is not None:
        rows = (await session.execute(
            select(ProjectionYearly).where(ProjectionYearly.mc_run_id == run.id).order_by(
                ProjectionYearly.year_idx))).scalars().all()
        yearly_rows = list(rows)

    rules = _flex_rules_for(scen)
    base = Decimal(str(scen.spending_gbp))
    horizon = scen.horizon_years
    by_year = {row.year_idx: row for row in yearly_rows}

    points: list[SpendingProfilePoint] = []
    ath = Decimal(str(scen.nw_seed_gbp))
    for year_idx in range(horizon):
        essential = _category_outflow_at(events, year_idx, "essential")
        discretionary = _category_outflow_at(events, year_idx, "discretionary")
        not_spending = _category_outflow_at(events, year_idx, "not_spending")

        # Income-shaped life events (positive delta) reduce the net spend
        # the user must sustain. We don't subtract them from any single
        # bucket — they net against `base` for the chart's footprint.
        ess_inflow = _category_inflow_at(events, year_idx, "essential")
        disc_inflow = _category_inflow_at(events, year_idx, "discretionary")
        net_base = base - ess_inflow - disc_inflow
        if net_base < 0:
            net_base = Decimal("0")

        flex_cut = Decimal("0")
        row = by_year.get(year_idx)
        if row is not None:
            ath = max(ath, Decimal(str(row.p50_portfolio_gbp)))
            flex_cut = _flex_cut_at_year(row, ath, rules, discretionary)

        total = net_base + essential + discretionary - flex_cut
        if total < 0:
            total = Decimal("0")

        points.append(
            SpendingProfilePoint(
                year_idx=year_idx,
                base_gbp=net_base,
                essential_gbp=essential,
                discretionary_gbp=discretionary,
                not_spending_gbp=not_spending,
                flex_cut_gbp=flex_cut,
                total_gbp=total,
            ))

    return SpendingProfileResponse(
        scenario_id=scenario_id,
        horizon_years=horizon,
        points=points,
    )


@router.post("/spending-profile-preview", response_model=SpendingProfileResponse)
async def preview_spending_profile(
    req: SpendingProfilePreviewRequest,
) -> SpendingProfileResponse:
    """Stateless variant — used by the What-If page where the user is
    tweaking in-memory life events that aren't persisted yet."""
    points: list[SpendingProfilePoint] = []
    base = Decimal(str(req.base_spending_gbp))
    # Flex rules are accepted for forward-compat but not applied here:
    # without a portfolio path we can't decide drawdown depth. The
    # paired /simulate call gives the exact cut; this endpoint just
    # gives the pre-cut profile for the stacked-area chart.
    _ = req.flex_rules

    for year_idx in range(req.horizon_years):
        essential = Decimal("0")
        discretionary = Decimal("0")
        not_spending = Decimal("0")
        ess_inflow = Decimal("0")
        disc_inflow = Decimal("0")
        for ev in req.life_events:
            if not ev.enabled:
                continue
            if year_idx < ev.year_start:
                continue
            end = ev.year_end if ev.year_end is not None else ev.year_start
            if year_idx > end:
                continue
            delta = Decimal(str(ev.delta_gbp_per_year or 0))
            if ev.category == "not_spending":
                not_spending += abs(delta)
                continue
            if delta < 0:
                if ev.category == "essential":
                    essential += -delta
                elif ev.category == "discretionary":
                    discretionary += -delta
            elif delta > 0:
                if ev.category == "essential":
                    ess_inflow += delta
                elif ev.category == "discretionary":
                    disc_inflow += delta
            if year_idx == ev.year_start and ev.one_time_amount_gbp is not None:
                ot = Decimal(str(ev.one_time_amount_gbp))
                if ot < 0:
                    if ev.category == "essential":
                        essential += -ot
                    elif ev.category == "discretionary":
                        discretionary += -ot
        net_base = base - ess_inflow - disc_inflow
        if net_base < 0:
            net_base = Decimal("0")
        total = net_base + essential + discretionary
        points.append(
            SpendingProfilePoint(
                year_idx=year_idx,
                base_gbp=net_base,
                essential_gbp=essential,
                discretionary_gbp=discretionary,
                not_spending_gbp=not_spending,
                flex_cut_gbp=Decimal("0"),
                total_gbp=total,
            ))
    return SpendingProfileResponse(
        scenario_id=None,
        horizon_years=req.horizon_years,
        points=points,
    )