"""Tests for the flex-spending engine."""
from __future__ import annotations

import numpy as np
import pytest

from fire_planner.flex_spending import FlexRule, applicable_cut, cuts_per_year


def test_applicable_cut_picks_deepest_rule() -> None:
    rules = [
        FlexRule(from_ath_pct=0.10, cut_discretionary_pct=0.20),
        FlexRule(from_ath_pct=0.30, cut_discretionary_pct=0.60),
        FlexRule(from_ath_pct=0.50, cut_discretionary_pct=0.90),
    ]
    # No drawdown — no cut.
    assert applicable_cut(0.0, rules) == 0.0
    # 9% drop — below first threshold.
    assert applicable_cut(0.09, rules) == 0.0
    # 15% drop — only first rule fires.
    assert applicable_cut(0.15, rules) == pytest.approx(0.20)
    # 35% drop — first + second; deepest cut wins (0.60, not 0.80).
    assert applicable_cut(0.35, rules) == pytest.approx(0.60)
    # 60% drop — all three; 0.90 wins.
    assert applicable_cut(0.60, rules) == pytest.approx(0.90)


def test_applicable_cut_empty_rules() -> None:
    assert applicable_cut(0.5, []) == 0.0


def test_cuts_per_year_handles_running_ath() -> None:
    # Single path. Year 0 seed=1000, year 1 = 1200 (new ATH), year 2 = 800
    # (-33% from ATH 1200), year 3 = 900 (still -25% from ATH 1200), year
    # 4 = 1300 (new ATH).
    portfolio = np.array([[1000, 1200, 800, 900, 1300]], dtype=np.float64)
    rules = [
        FlexRule(from_ath_pct=0.10, cut_discretionary_pct=0.20),
        FlexRule(from_ath_pct=0.30, cut_discretionary_pct=0.60),
    ]
    cuts = cuts_per_year(portfolio, rules)
    # cuts[:, y] uses portfolio[:, y] (start-of-year decision based on
    # the prior year's close).
    # y=0: portfolio=1000 == ATH → 0
    # y=1: portfolio=1200 == ATH → 0
    # y=2: drawdown = 1 - 800/1200 = 0.333 → 0.60
    # y=3: drawdown = 1 - 900/1200 = 0.25 → 0.20
    assert cuts.shape == (1, 4)
    assert cuts[0, 0] == pytest.approx(0.0)
    assert cuts[0, 1] == pytest.approx(0.0)
    assert cuts[0, 2] == pytest.approx(0.60)
    assert cuts[0, 3] == pytest.approx(0.20)


def test_cuts_per_year_no_rules_returns_zeros() -> None:
    portfolio = np.array([[1000, 800, 600]], dtype=np.float64)
    cuts = cuts_per_year(portfolio, [])
    assert cuts.shape == (1, 2)
    assert (cuts == 0).all()