"""ProjectionLab-style flex-spending rules.

A `FlexRule` says "if the portfolio is at least ``from_ath_pct`` below its
running all-time-high, trim discretionary spending by ``cut_discretionary_pct``".
Multiple rules stack via "deepest applicable cut wins" — users specify
*cumulative* cuts at each tier, so a [-0.10 → 20%, -0.30 → 60%] config
trims by 60% (not 80%) at -30%.

The engine path:

  per year y, per path p:
    drawdown[p,y] = 1 - portfolio[p,y] / ath[p,y]
    cut_pct[p,y] = max(rule.cut for rule in flex_rules if drawdown[p,y] >= rule.from_ath_pct)
    discretionary_after_flex[p,y] = discretionary_baseline[y] * (1 - cut_pct[p,y])

The cuts are applied to the *baseline* discretionary spend each year (so a
£10k/y travel budget cut by 60% becomes £4k that year), and the saved
amount is *not* drawn from the portfolio. The simulator subtracts the
saved amount from the cashflow drawdown before calling the strategy.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import numpy.typing as npt


@dataclass(frozen=True)
class FlexRule:
    """Engine-level flex rule. ``from_ath_pct`` is the absolute drop
    magnitude (positive fraction); ``cut_discretionary_pct`` is the
    fraction to remove from discretionary spending at that depth."""
    from_ath_pct: float
    cut_discretionary_pct: float


def applicable_cut(drawdown: float, rules: list[FlexRule]) -> float:
    """Return the cut fraction for a single (path, year) pair.

    ``drawdown`` is 1 − portfolio/ath (in the [0, 1] range — clamp inside
    the simulator before calling). The deepest rule whose threshold is
    satisfied wins.
    """
    if not rules:
        return 0.0
    best = 0.0
    for rule in rules:
        if drawdown >= rule.from_ath_pct and rule.cut_discretionary_pct > best:
            best = rule.cut_discretionary_pct
    return best


def cuts_per_year(
    portfolio_real: npt.NDArray[np.float64],
    rules: list[FlexRule],
) -> npt.NDArray[np.float64]:
    """Vectorised version of ``applicable_cut`` across every (path, year).

    ``portfolio_real`` shape: ``(n_paths, n_years + 1)`` — index 0 is the
    seed, last column is the horizon. Returns ``(n_paths, n_years)``: the
    cut applied at the start of year ``y`` is decided by the portfolio
    *after year y-1's close* (i.e. column ``y`` in the input).
    """
    if not rules or portfolio_real.size == 0:
        return np.zeros((portfolio_real.shape[0], portfolio_real.shape[1] - 1),
                        dtype=np.float64)
    n_paths, ncols = portfolio_real.shape
    n_years = ncols - 1
    # Running ATH per path. np.maximum.accumulate over axis=1 gives us
    # the running max — exactly what we want.
    ath = np.maximum.accumulate(portfolio_real, axis=1)
    # Avoid divide-by-zero. If ATH is 0 (only happens if seed is 0 and the
    # portfolio never grew), drawdown is treated as 0.
    safe_ath = np.where(ath > 0, ath, 1.0)
    drawdown = np.clip(1.0 - portfolio_real / safe_ath, 0.0, 1.0)
    cuts = np.zeros((n_paths, n_years), dtype=np.float64)
    sorted_rules = sorted(rules,
                          key=lambda r: r.cut_discretionary_pct,
                          reverse=True)
    for rule in sorted_rules:
        # Each rule's cut applies wherever drawdown >= threshold AND a
        # higher cut hasn't already been recorded (because we iterate
        # rules from largest cut down).
        # Drawdown at year y end-of-year-(y-1) — column y of drawdown.
        mask = (drawdown[:, :n_years] >= rule.from_ath_pct) & (
            cuts < rule.cut_discretionary_pct)
        cuts[mask] = rule.cut_discretionary_pct
    return cuts