"""VPW — Variable Percentage Withdrawal (Bogleheads).

Withdrawal rate is the standard PMT (annuity-payment) formula given a
target real return and the years remaining:

    rate(n, r) = r / (1 - (1 + r)^-n)

At year `y` of an `H`-year horizon, withdraw
`portfolio * rate(H - y, expected_real_return)`. The withdrawal scales
with the portfolio — bear markets cut spending immediately, bull
markets allow more — eliminating ruin risk at the cost of variable
income.

Bogleheads VPW table values (60% stocks, 40% bonds, 5% real expected):
- Age 35, 60y horizon: 5.30%
- Age 50, 45y horizon: 5.86%
- Age 65, 30y horizon: 7.09%
- Age 80, 15y horizon: 11.42%

Default `expected_real_return=0.05` matches Bogleheads' 60/40 assumption.
"""
from __future__ import annotations

from fire_planner.strategies.base import StrategyState, WithdrawalStrategy

DEFAULT_EXPECTED_REAL_RETURN = 0.05


def pmt_rate(years_remaining: int, real_rate: float) -> float:
    """PMT formula, capped at 1.0.

    Ordinary-annuity convention matches Bogleheads' VPW table for
    horizons > 1y. At 1y remaining the textbook formula returns
    `(1+r)` because of end-of-period interest accrual, which would
    propose a withdrawal larger than the portfolio — Bogleheads caps
    at 100% in that case, and so do we.
    """
    if years_remaining <= 0:
        return 1.0
    if abs(real_rate) < 1e-9:
        return 1.0 / years_remaining
    return min(1.0, real_rate / (1.0 - (1.0 + real_rate)**-years_remaining))


class VpwStrategy(WithdrawalStrategy):
    name = "vpw"

    def __init__(self, expected_real_return: float = DEFAULT_EXPECTED_REAL_RETURN) -> None:
        self.expected_real_return = expected_real_return

    def propose_withdrawal(self, state: StrategyState) -> float:
        if state.portfolio <= 0:
            return 0.0
        years_left = state.horizon_years - state.year_idx
        rate = pmt_rate(years_left, self.expected_real_return)
        return state.portfolio * rate


class VpwWithFloorStrategy(WithdrawalStrategy):
    """VPW with a real-GBP floor — the never-drop-below safety net.

    Each year propose `max(floor, vpw_proposed)`, then clip to portfolio
    so we cannot withdraw more than exists. The floor is the binding
    constraint in bad sequences; in good sequences VPW dominates and
    spending scales up. The simulator's success_mask uses the
    portfolio-positive-through-interim-years check, so a floor that
    drains the portfolio early is still penalised the right way.
    """
    name = "vpw_floor"

    def __init__(self,
                 floor: float,
                 expected_real_return: float = DEFAULT_EXPECTED_REAL_RETURN) -> None:
        self.floor = floor
        self.expected_real_return = expected_real_return

    def propose_withdrawal(self, state: StrategyState) -> float:
        if state.portfolio <= 0:
            return 0.0
        years_left = state.horizon_years - state.year_idx
        rate = pmt_rate(years_left, self.expected_real_return)
        vpw_proposed = state.portfolio * rate
        return min(state.portfolio, max(self.floor, vpw_proposed))