stack/tests/lead/test_benchmark.py

"""lead.benchmark reproduces the §3 worked examples: Box 3.3.2, Table 3.3.3,
Table 3.4.2, Tables 3.4.3.1–3.4.3.9, §3.5.1–3.5.2, Table 3.5.3.2 and
Table 3.5.4."""

from __future__ import annotations

import pytest

from lead.benchmark import (
    CAPPED,
    CAPPED_AT_UPPER,
    FLOORED,
    RAISED_TO_LOWER,
    WITHIN_CAP,
    WITHIN_GUARDRAIL,
    base_year_alignment_minimum,
    base_year_is_sufficient,
    base_year_weights,
    benchmark_adjustment_cap,
    benchmark_adjustment_cap_rate,
    benchmark_dollars,
    capped_risk_score,
    contributes_to_voluntary_benchmark,
    cumulative_acpt,
    expenditure_pbpm,
    growth_rate,
    historical_benchmark,
    hybrid_adjustment_factor,
    month_weighted_average,
    national_trend_weight,
    performance_year_benchmark_pbpm,
    prior_savings_adjustment,
    prior_savings_cap,
    prior_savings_pbpm,
    ratio,
    regional_efficiency_adjustment,
    regional_efficiency_eligible,
    regional_ffs_expenditure,
    risk_standardize,
    spending_designation,
    three_way_blended_trend,
    total_benchmark,
    two_way_blended_trend,
    voluntary_benchmark_minimum_met,
    voluntary_guardrail,
)
from lead.money import cents
from tests.lead.paper import (
    REA_TABLES,
    TABLE_3_3_3,
    TABLE_3_4_2_DIFFERENCE,
    TABLE_3_5_3_2,
    TABLE_3_5_4_CLAIMS,
    TABLE_3_5_4_TOTALS,
    TABLE_3_5_4_VOLUNTARY,
)

# ── Box 3.3.2 and Table 3.3.3 ─────────────────────────────────────────────


def test_box_3_3_2_national_regional_blended_trend():
    national = growth_rate(1_120, 1_000)
    regional = growth_rate(1_155, 1_050)
    weight = national_trend_weight([(20_000, 100_000), (5_000, 50_000)])
    assert (national, round(regional, 3), round(weight, 6)) == (1.12, 1.1, 0.18)
    trend = two_way_blended_trend(national, regional, weight)
    assert round(trend, 3) == 1.104  # printed 1.104; the text's "1.103" is a slip


def test_trend_inputs_without_positive_denominators_are_missing():
    assert growth_rate(0, 1_000) is None
    assert growth_rate(1_000, 0) is None
    assert national_trend_weight([]) is None
    assert national_trend_weight([(10, 0)]) is None
    assert two_way_blended_trend(None, 1.1, 0.2) is None
    assert ratio(None, 1.0) is None


@pytest.mark.parametrize("year", ["BY1", "BY2", "BY3"])
def test_table_3_3_3_lines_3_to_9(year):
    payments, months, risk, line3, line5, trend, line7, weight, line9 = TABLE_3_3_3[
        year
    ]
    pbpm = expenditure_pbpm(payments, months)
    standardized = risk_standardize(pbpm, risk)
    # PBPMs are printed to the cent but carried unrounded: line 5 is
    # $857.03 from the unrounded line 3, $857.02 from the printed $961.58.
    assert cents(pbpm) == line3
    assert cents(standardized) == line5
    assert cents(standardized * trend) == line7
    assert cents(standardized * trend * weight) == line9


def test_table_3_3_3_historical_benchmark():
    years = ("BY1", "BY2", "BY3")
    standardized = {
        y: risk_standardize(expenditure_pbpm(*TABLE_3_3_3[y][:2]), TABLE_3_3_3[y][2])
        for y in years
    }
    trends = {y: TABLE_3_3_3[y][5] for y in years}
    weights = base_year_weights("Newly Entering", years)
    benchmark = historical_benchmark(standardized, trends, weights)
    assert cents(benchmark) == TABLE_3_3_3["historical_benchmark"]
    # The printed line 9 cents sum to 943.63: the total is unrounded too.
    assert sum(TABLE_3_3_3[y][8] for y in years) == pytest.approx(943.63)


def test_historical_benchmark_needs_every_weighted_year():
    assert historical_benchmark({"BY3": 900.0}, {"BY3": 1.0}, {}) is None
    assert historical_benchmark({}, {"BY3": 1.0}, {"BY3": 1.0}) is None
    assert historical_benchmark({"BY3": 900.0}, {"BY3": 1.0}, {"BY3": 1.0}) == 900.0


def test_risk_standardize_needs_a_positive_score():
    assert risk_standardize(1_000.0, 0.0) is None
    assert expenditure_pbpm(1_000.0, 0) is None


# ── §3.3.3–3.3.4 base-year weights and sufficiency ────────────────────────


@pytest.mark.parametrize(
    "entry,years,weights",
    [
        ("Newly Entering", ("BY1", "BY2", "BY3"), {"BY1": 0.1, "BY2": 0.3, "BY3": 0.6}),
        ("Renewing", ("BY1", "BY2", "BY3"), {"BY1": 1 / 3, "BY2": 1 / 3, "BY3": 1 / 3}),
        ("Newly Entering", ("BY1", "BY3"), {"BY1": 1 / 3, "BY3": 2 / 3}),
        ("Renewing", ("BY2", "BY3"), {"BY2": 0.5, "BY3": 0.5}),
        ("Renewing", ("BY2",), {"BY2": 1.0}),
        ("Newly Entering", (), {}),
    ],
)
def test_base_year_weights(entry, years, weights):
    assert base_year_weights(entry, years) == weights


def test_base_year_weights_refuse_unknown_labels():
    with pytest.raises(ValueError, match="unknown entry type"):
        base_year_weights("New", ("BY3",))
    with pytest.raises(ValueError, match="unknown base years"):
        base_year_weights("Renewing", ("BY4",))


def test_table_2_6_1_base_year_minimums():
    assert base_year_alignment_minimum("standard", 2027) == 3_000
    assert base_year_alignment_minimum("newly_entering", 2029) == 1_800
    assert base_year_alignment_minimum("high_needs_esrd", 2030) == 825
    assert base_year_alignment_minimum("high_needs_esrd", 2036) == 1_000
    assert base_year_is_sufficient(600, "newly_entering", 2027)
    assert not base_year_is_sufficient(599, "newly_entering", 2027)


# ── §3.2.5 voluntary benchmarks ───────────────────────────────────────────


@pytest.mark.parametrize(
    "voluntary,expected",
    [
        (1_150.0, (1_100.0, CAPPED_AT_UPPER)),
        (850.0, (900.0, RAISED_TO_LOWER)),
        (1_050.0, (1_050.0, WITHIN_GUARDRAIL)),
    ],
)
def test_voluntary_guardrail_example(voluntary, expected):
    benchmark, status = voluntary_guardrail(voluntary, 1_000.0)
    assert (round(benchmark, 6), status) == expected


def test_voluntary_contribution_and_minimum():
    assert contributes_to_voluntary_benchmark("A&D", None)
    assert contributes_to_voluntary_benchmark("High Needs", "High Needs")
    assert not contributes_to_voluntary_benchmark("High Needs", "A&D")
    assert not contributes_to_voluntary_benchmark("High Needs", None)
    assert voluntary_benchmark_minimum_met(500)
    assert not voluntary_benchmark_minimum_met(499)


# ── §3.4.1–3.4.2 regional FFS and spending designation ────────────────────


def test_regional_ffs_expenditure_weights_risk_standardized_counties():
    # County A: $1,100 / 1.10 = $1,000 at 75%; county B: $1,200 / 1.00 at 25%.
    assert regional_ffs_expenditure(
        [(1_100.0, 1.10, 750.0), (1_200.0, 1.0, 250.0)]
    ) == pytest.approx(1_050.0)
    assert regional_ffs_expenditure([(1_100.0, None, 750.0)]) is None
    assert regional_ffs_expenditure([(1_100.0, 1.1, 0.0)]) is None


def test_table_3_4_2_spending_designation():
    rows = [(r[0], r[1], r[2]) for r in REA_TABLES.values()]
    result = spending_designation(rows)
    assert result.weighted_difference == pytest.approx(TABLE_3_4_2_DIFFERENCE)
    assert result.designation == "Lower-spending"
    assert (
        spending_designation([(900.0, 1_000.0, 1.0)]).designation == "Higher-spending"
    )
    # The paper names no tie.
    assert spending_designation([(1_000.0, 1_000.0, 1.0)]).designation is None
    with pytest.raises(ValueError, match="BY 3 aligned months"):
        spending_designation([(1_000.0, 900.0, 0.0)])


# ── §3.4.3 Regional Efficiency Adjustment ─────────────────────────────────


@pytest.mark.parametrize("category", list(REA_TABLES))
def test_tables_3_4_3_1_to_3_4_3_4(category):
    regional, historical, _, uncapped, uspcc, risk, cap, applied, adjusted = REA_TABLES[
        category
    ]
    rea = regional_efficiency_adjustment(
        regional,
        historical,
        eligible=regional_efficiency_eligible("Lower-spending", "Global"),
        uspcc_amount=uspcc,
        py_average_risk_score=risk,
        cap_rate=benchmark_adjustment_cap_rate(reduced=False),
    )
    assert rea.uncapped == uncapped
    assert cents(rea.cap) == cap
    assert cents(rea.applied) == applied
    assert cents(historical + rea.applied) == adjusted


def test_rea_eligibility_and_missing_cap_inputs():
    assert not regional_efficiency_eligible("Higher-spending", "Global")
    assert not regional_efficiency_eligible("Lower-spending", "Professional")
    kwargs = dict(uspcc_amount=None, py_average_risk_score=1.0, cap_rate=0.05)
    # A positive REA never falls back to its uncapped amount.
    assert (
        regional_efficiency_adjustment(1_050, 1_000, eligible=True, **kwargs).applied
        is None
    )
    assert (
        regional_efficiency_adjustment(950, 1_000, eligible=True, **kwargs).applied
        == 0.0
    )
    assert (
        regional_efficiency_adjustment(1_050, 1_000, eligible=False, **kwargs).applied
        == 0.0
    )


def test_reduced_cap_rate():
    assert benchmark_adjustment_cap_rate(reduced=True) == 0.03
    assert benchmark_adjustment_cap(1_000, 1.02, 0.03) == pytest.approx(30.6)


def test_table_3_4_3_6_weighted_rea_with_an_excluded_tin():
    rea = regional_efficiency_adjustment(
        3_700,
        3_000,
        eligible=True,
        uspcc_amount=3_500,
        py_average_risk_score=1.03,
        cap_rate=0.05,
    )
    assert (rea.difference, rea.uncapped, cents(rea.applied)) == (700, 350, 180.25)
    assert cents(month_weighted_average([(rea.applied, 8_000), (0.0, 2_000)])) == 144.20


# ── §3.4.4 Prior Savings Adjustment ───────────────────────────────────────


def test_table_3_4_3_7_and_example_1():
    rows = [
        (120_000_000, 114_250_000, 115_000),
        (125_000_000, 126_200_000, 120_000),
        (130_000_000, 121_250_000, 125_000),
    ]
    pbpm = [prior_savings_pbpm(*row) for row in rows]
    assert pbpm == [50.0, -10.0, 70.0]
    average = sum(pbpm) / 3
    assert cents(average) == 36.67
    psa = prior_savings_adjustment(average, 120_000, 240_000, cap=None)
    assert psa.eligible and psa.proration_factor == 0.5
    assert cents(psa.prorated_average_savings) == 18.33
    assert cents(psa.uncapped) == 9.17
    assert psa.applied is None  # eligible, but no cap supplied


def test_table_3_4_3_8_9_examples_2_and_3():
    average = month_weighted_average([(30.0, 80_000), (70.0, 20_000)])
    assert average == pytest.approx(38.0)
    cap = prior_savings_cap(
        [(0.80, 1_000, 1.02), (0.05, 8_000, 0.99), (0.15, 3_500, 1.03)], cap_rate=0.05
    )
    assert cents(cap) == 87.64
    psa = prior_savings_adjustment(average, 100_000, 125_000, cap)
    assert psa.proration_factor == 0.8
    assert cents(psa.prorated_average_savings) == 30.40
    assert cents(psa.applied) == 15.20


def test_prior_savings_eligibility_and_proration_cap():
    not_eligible = prior_savings_adjustment(-5.0, 100_000, 100_000, cap=50.0)
    assert (not_eligible.eligible, not_eligible.applied) == (False, 0.0)
    shrank = prior_savings_adjustment(40.0, 200_000, 100_000, cap=5.0)
    assert shrank.proration_factor == 1.0
    assert shrank.applied == 5.0  # capped
    with pytest.raises(ValueError, match="PY beneficiary months"):
        prior_savings_adjustment(40.0, 1, 0, cap=None)
    assert prior_savings_pbpm(1.0, 0.0, 0) is None


# ── §3.5.1 three-way trend and ACPT ───────────────────────────────────────


@pytest.mark.parametrize(
    "year,acpt,applied",
    [
        (2027, 1.070, 1.053),  # raw 1.0567 capped at two-way + 0.3 points
        (2027, 1.020, 1.048),  # raw 1.04 raised to two-way − 0.2 points
        (2027, 1.055, (2 / 3) * 1.05 + 1.055 / 3),  # within the band
        (2028, 1.070, 1.056),  # PY 2 caps at +0.6 points
        (2028, 1.065, 1.055),  # within PY 2's wider band
        (2032, 0.950, 1.04),  # PY 6: −1.0 points
    ],
)
def test_three_way_blended_trend_guardrail(year, acpt, applied):
    trend = three_way_blended_trend(1.05, acpt, year)
    assert trend.applied == pytest.approx(applied)
    assert trend.lower_bound <= trend.applied <= trend.upper_bound


def test_cumulative_acpt():
    assert cumulative_acpt([0.04, 0.05]) == pytest.approx(1.092)
    assert cumulative_acpt([0.04, None]) is None
    assert cumulative_acpt([]) is None


# ── §3.5.2 risk adjustment and the growth cap ─────────────────────────────


def test_section_3_5_2_examples():
    assert performance_year_benchmark_pbpm(1_100.0, 0.0, 1.0, 1.02) == pytest.approx(
        1_122.0
    )
    score, status = capped_risk_score(1.000, 1.050, "A&D")
    assert (score, status) == (1.03, CAPPED)
    assert performance_year_benchmark_pbpm(1_100.0, 0.0, 1.0, score) == pytest.approx(
        1_133.0
    )


def test_risk_score_cap_by_category():
    assert capped_risk_score(1.0, 0.90, "ESRD") == (0.97, FLOORED)
    assert capped_risk_score(1.0, 1.035, "High Needs") == (1.035, WITHIN_CAP)
    assert capped_risk_score(1.0, 1.05, "High Needs") == (1.04, CAPPED)


def test_missing_adjustment_stays_missing():
    assert performance_year_benchmark_pbpm(1_000.0, None, 1.03, 1.0) is None


# ── §3.5.3 hybrid alignment and §3.5.4 combination ────────────────────────


def test_table_3_5_3_2_hybrid_alignment():
    t = TABLE_3_5_3_2
    assert cents(t["1"] * t["4"]) == t["5"]
    dollars = benchmark_dollars(t["1"], t["2"], [(t["3"], t["4"])])
    assert (dollars.full_year, dollars.hybrid, dollars.total) == (
        t["6"],
        t["7"],
        t["8"],
    )
    assert hybrid_adjustment_factor(1_009.9, 1_000.0) == pytest.approx(1.0099)
    assert hybrid_adjustment_factor(1.0, 0.0) is None


@pytest.mark.parametrize("category", list(TABLE_3_5_4_CLAIMS))
def test_table_3_5_4_claims_line_8(category):
    pbpm, trend, risk, full, hybrid, factor, line8 = TABLE_3_5_4_CLAIMS[category]
    py_pbpm = performance_year_benchmark_pbpm(pbpm, 0.0, trend, risk)
    assert benchmark_dollars(py_pbpm, full, [(hybrid, factor)]).total == line8


@pytest.mark.parametrize("category", list(TABLE_3_5_4_VOLUNTARY))
def test_table_3_5_4_voluntary_line_22(category):
    pbpm, trend, risk, full, months, factors, line22 = TABLE_3_5_4_VOLUNTARY[category]
    py_pbpm = performance_year_benchmark_pbpm(pbpm, 0.0, trend, risk)
    hybrid = list(zip(months, factors, strict=True))
    assert benchmark_dollars(py_pbpm, full, hybrid).total == line22


def test_table_3_5_4_totals_sum_the_cent_rounded_cells():
    claims = {c: v[-1] for c, v in TABLE_3_5_4_CLAIMS.items()}
    voluntary = {c: v[-1] for c, v in TABLE_3_5_4_VOLUNTARY.items()}
    assert total_benchmark(claims.values()) == TABLE_3_5_4_TOTALS["8"]
    assert total_benchmark(voluntary.values()) == TABLE_3_5_4_TOTALS["22"]
    line24 = {c: total_benchmark([claims[c], voluntary[c]]) for c in claims}
    for category, expected in TABLE_3_5_4_TOTALS["24"].items():
        if category != "TOTAL":
            assert line24[category] == expected
    assert total_benchmark(line24.values()) == TABLE_3_5_4_TOTALS["24"]["TOTAL"]


def test_table_3_5_4_esrd_line_24_needs_cent_rounded_cells():
    # Unrounded, ESRD line 8 + line 22 = 17,675,378.2467 → .25, not the printed .24.
    pbpm, trend, risk, full, hybrid, factor, _ = TABLE_3_5_4_CLAIMS["ESRD"]
    claims = pbpm * trend * risk * (full + hybrid * factor)
    pbpm, trend, risk, full, months, factors, _ = TABLE_3_5_4_VOLUNTARY["ESRD"]
    voluntary = (
        pbpm
        * trend
        * risk
        * (full + sum(m * f for m, f in zip(months, factors, strict=True)))
    )
    assert cents(claims + voluntary) == 17_675_378.25
    assert TABLE_3_5_4_TOTALS["24"]["ESRD"] == 17_675_378.24