fire-planner/fire_planner/api/cashflow.py

"""Cashflow Sankey data for one year — ProjectionLab parity Wave 1.C.2.

Returns a ``{sources, sinks}`` map of named tributaries / drains. Sums
conserve at the Sankey level: total inflow == total outflow, with
"net change in NW" absorbing the residual so the diagram balances.
"""
from __future__ import annotations

from decimal import Decimal

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

from fire_planner.api.dependencies import get_session
from fire_planner.api.schemas import CashflowResponse
from fire_planner.db import IncomeStream, LifeEvent, McRun, ProjectionYearly, Scenario

router = APIRouter(prefix="/scenarios", tags=["cashflow"])


def _income_inflow_at(streams: list[IncomeStream], year_idx: int) -> dict[str, Decimal]:
    """Per-stream-kind inflow at ``year_idx``."""
    out: dict[str, Decimal] = {}
    for s in streams:
        if not s.enabled:
            continue
        if year_idx < s.start_year:
            continue
        if s.end_year is not None and year_idx > s.end_year:
            continue
        years_active = year_idx - s.start_year
        growth = (Decimal("1") + Decimal(str(s.growth_pct)))**years_active
        amount = Decimal(str(s.amount_gbp_per_year)) * growth
        key = f"income:{s.kind}"
        out[key] = out.get(key, Decimal("0")) + amount
    return out


def _life_event_inflow_at(events: list[LifeEvent], year_idx: int) -> dict[str, Decimal]:
    out: dict[str, Decimal] = {}
    for ev in 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.year_start == year_idx and ev.one_time_amount_gbp is not None:
            delta += Decimal(str(ev.one_time_amount_gbp))
        if delta == 0:
            continue
        key = f"event:{ev.name}"
        if delta > 0:
            out[key] = out.get(key, Decimal("0")) + delta
    return out


@router.get("/{scenario_id}/cashflow", response_model=CashflowResponse)
async def get_cashflow(
    scenario_id: int,
    year: int = Query(ge=0, le=100, description="Year offset from scenario start (0 = now)"),
    session: AsyncSession = Depends(get_session),
) -> CashflowResponse:
    scen = await session.get(Scenario, scenario_id)
    if scen is None:
        raise HTTPException(status_code=404, detail="Scenario not found")
    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 None:
        raise HTTPException(status_code=404, detail="No MC runs persisted for this scenario yet")

    yearly_rows = (await session.execute(
        select(ProjectionYearly).where(ProjectionYearly.mc_run_id == run.id).order_by(
            ProjectionYearly.year_idx))).scalars().all()
    by_year = {r.year_idx: r for r in yearly_rows}
    if year not in by_year:
        raise HTTPException(status_code=404, detail=f"No projection row for year {year}")

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

    row = by_year[year]
    prev = by_year.get(year - 1) if year > 0 else None
    nw = Decimal(str(row.p50_portfolio_gbp))
    prev_nw = (Decimal(str(prev.p50_portfolio_gbp)) if prev is not None else Decimal(str(
        scen.nw_seed_gbp)))
    nw_change = nw - prev_nw

    sources: dict[str, Decimal] = {}
    sources.update(_income_inflow_at(streams, year))
    sources.update(_life_event_inflow_at(events, year))
    if scen.savings_per_year_gbp > 0:
        sources["savings"] = Decimal(str(scen.savings_per_year_gbp))

    spending = Decimal(str(row.p50_withdrawal_gbp))
    taxes = Decimal(str(row.p50_tax_gbp))
    sinks: dict[str, Decimal] = {
        "spending": spending,
        "taxes": taxes,
    }

    sources_total = sum(sources.values(), Decimal("0"))
    sinks_total = sum(sinks.values(), Decimal("0"))
    investment_returns = nw_change + sinks_total - sources_total
    if investment_returns > 0:
        sources["investment_returns"] = investment_returns
        sources_total += investment_returns
    elif investment_returns < 0:
        sinks["portfolio_drawdown"] = -investment_returns
        sinks_total += -investment_returns

    delta = sources_total - sinks_total
    if delta > 0:
        sinks["net_change_in_nw"] = delta
    elif delta < 0:
        sources["net_change_in_nw"] = -delta

    return CashflowResponse(
        scenario_id=scenario_id,
        year=year,
        sources=sources,
        sinks=sinks,
    )