stack/tests/lead/test_express_benchmark.py

"""lead.express.benchmark: the §3 worked examples as frames on polars and
DuckDB, plus nulls for missing inputs."""

from __future__ import annotations

import narwhals as nw
import polars as pl
import pytest

from lead import benchmark as scalar
from lead.express import cents
from lead.express.benchmark import (
    RISK_CAP_REFERENCE_MISSING,
    benchmark_adjustment_cap,
    benchmark_dollars,
    capped_risk_score,
    contributes_to_voluntary_benchmark,
    county_market_share,
    county_weight,
    growth_rate,
    month_weighted_mean,
    performance_year_benchmark_pbpm,
    regional_efficiency_adjustment,
    regional_efficiency_eligible,
    risk_growth_cap_status,
    risk_standardized_pbpm,
    three_way_blended_trend,
    trended_weighted_pbpm,
    two_way_blended_trend,
    uncapped_regional_efficiency_adjustment,
    voluntary_guardrail,
    voluntary_guardrail_status,
)
from lead.money import cents as scalar_cents
from tests.lead.paper import (
    REA_TABLES,
    TABLE_3_3_3,
    TABLE_3_5_4_CLAIMS,
    TABLE_3_5_4_TOTALS,
    TABLE_3_5_4_VOLUNTARY,
)


def test_box_3_3_2_as_a_county_frame(backend):
    counties = {
        "aligned": [20_000.0, 5_000.0],
        "eligible": [100_000.0, 50_000.0],
        "total": [25_000.0, 25_000.0],
    }
    rows = backend.select(
        counties,
        contribution=county_market_share("aligned", "eligible")
        * county_weight("aligned", "total"),
    )
    weight = sum(r["contribution"] for r in rows)
    assert weight == pytest.approx(0.18)
    (trend,) = backend.select(
        {
            "nat_by3": [1_120.0],
            "nat_by1": [1_000.0],
            "reg_by3": [1_155.0],
            "reg_by1": [1_050.0],
            "w": [weight],
        },
        trend=two_way_blended_trend(
            growth_rate("nat_by3", "nat_by1"), growth_rate("reg_by3", "reg_by1"), "w"
        ),
    )
    assert round(trend["trend"], 3) == 1.104


def test_growth_and_shares_are_null_without_positive_denominators(backend):
    data = {"a": [0.0, 10.0, None], "b": [5.0, 0.0, 5.0]}
    rows = backend.select(
        data,
        {"a": pl.Float64, "b": pl.Float64},
        g=growth_rate("a", "b"),
        s=county_market_share("a", "b"),
    )
    assert [(r["g"], r["s"]) for r in rows] == [(None, 0.0), (None, None), (None, None)]


def test_table_3_3_3_as_a_base_year_frame(backend):
    years = ("BY1", "BY2", "BY3")
    data = {
        "year": list(years),
        "paid": [TABLE_3_3_3[y][0] for y in years],
        "months": [float(TABLE_3_3_3[y][1]) for y in years],
        "risk": [TABLE_3_3_3[y][2] for y in years],
        "trend": [TABLE_3_3_3[y][5] for y in years],
        "weight": [TABLE_3_3_3[y][7] for y in years],
    }
    rs = risk_standardized_pbpm("paid", "months", "risk")
    rows = backend.select(
        data,
        year=nw.col("year"),
        line5=cents(rs),
        line9=cents(trended_weighted_pbpm(rs, "trend", "weight")),
        raw=trended_weighted_pbpm(rs, "trend", "weight"),
    )
    by_year = {r["year"]: r for r in rows}
    for year in years:
        assert by_year[year]["line5"] == TABLE_3_3_3[year][4]
        assert by_year[year]["line9"] == TABLE_3_3_3[year][8]
    assert (
        scalar_cents(sum(r["raw"] for r in rows)) == TABLE_3_3_3["historical_benchmark"]
    )


def test_rea_tables_as_a_category_frame(backend):
    data = {
        "category": list(REA_TABLES),
        "regional": [float(v[0]) for v in REA_TABLES.values()],
        "historical": [float(v[1]) for v in REA_TABLES.values()],
        "uspcc": [float(v[4]) for v in REA_TABLES.values()],
        "risk": [v[5] for v in REA_TABLES.values()],
        "designation": ["Lower-spending"] * 3,
        "option": ["Global"] * 3,
    }
    rea = regional_efficiency_adjustment(
        "regional",
        "historical",
        eligible=regional_efficiency_eligible("designation", "option"),
        uspcc_amount="uspcc",
        py_average_risk_score="risk",
        cap_rate=0.05,
    )
    rows = backend.select(
        data,
        uncapped=uncapped_regional_efficiency_adjustment("regional", "historical"),
        cap=cents(benchmark_adjustment_cap("uspcc", "risk", 0.05)),
        applied=cents(rea),
        adjusted=cents(nw.col("historical") + rea),
    )
    for values, result in zip(REA_TABLES.values(), rows, strict=True):
        assert (
            result["uncapped"],
            result["cap"],
            result["applied"],
            result["adjusted"],
        ) == (values[3], values[6], values[7], values[8])


def test_rea_nulls_and_zeros(backend):
    data = {
        "regional": [1_050.0, 1_050.0, 950.0, 1_050.0, None],
        "historical": [1_000.0] * 5,
        "uspcc": [None, 1_000.0, None, 1_000.0, 1_000.0],
        "eligible": [True, None, True, False, True],
    }
    schema = {
        "regional": pl.Float64,
        "historical": pl.Float64,
        "uspcc": pl.Float64,
        "eligible": pl.Boolean,
    }
    rea = regional_efficiency_adjustment(
        "regional",
        "historical",
        eligible="eligible",
        uspcc_amount="uspcc",
        py_average_risk_score=nw.lit(1.0),
        cap_rate="cap_rate",
    )
    frame_rows = backend.select(
        {**data, "cap_rate": [0.05] * 5}, {**schema, "cap_rate": pl.Float64}, rea=rea
    )
    # missing cap on a positive REA, unknown eligibility, at/above regional,
    # ineligible, missing regional
    assert [r["rea"] for r in frame_rows] == [None, None, 0.0, 0.0, None]


def test_three_way_trend_per_row_performance_year(backend):
    data = {
        "two_way": [1.05, 1.05, 1.05, 1.05],
        "acpt": [1.07, 1.02, 1.07, 1.07],
        "py": [2027, 2027, 2028, 2040],
    }
    rows = backend.select(data, t=three_way_blended_trend("two_way", "acpt", "py"))
    expected = [
        scalar.three_way_blended_trend(1.05, a, y).applied
        for a, y in ((1.07, 2027), (1.02, 2027), (1.07, 2028))
    ]
    assert [r["t"] for r in rows] == [*expected, None]
    (literal,) = backend.select(
        {"two_way": [1.05], "acpt": [1.07]},
        t=three_way_blended_trend("two_way", "acpt", 2027),
    )
    assert literal["t"] == expected[0]


def test_capped_risk_score_and_status(backend):
    data = {
        "by3": [1.0, 1.0, 1.0, None, 1.0],
        "py": [1.05, 0.90, 1.035, 1.2, 1.2],
        "cat": ["A&D", "ESRD", "High Needs", "A&D", "Other"],
    }
    rows = backend.select(
        data,
        {"by3": pl.Float64, "py": pl.Float64, "cat": pl.String},
        score=capped_risk_score("by3", "py", "cat"),
        status=risk_growth_cap_status("by3", "py", "cat"),
    )
    expected = [
        scalar.capped_risk_score(1.0, py, cat)
        for py, cat in ((1.05, "A&D"), (0.90, "ESRD"), (1.035, "High Needs"))
    ]
    assert [(r["score"], r["status"]) for r in rows[:3]] == expected
    # No BY 3 reference or unknown category: no silent pass-through.
    assert [(r["score"], r["status"]) for r in rows[3:]] == [
        (None, RISK_CAP_REFERENCE_MISSING)
    ] * 2


def test_voluntary_guardrail_and_contribution(backend):
    data = {
        "vol": [1_150.0, 850.0, 1_050.0, 1_000.0],
        "claims": [1_000.0, 1_000.0, 1_000.0, None],
    }
    rows = backend.select(
        data,
        {"vol": pl.Float64, "claims": pl.Float64},
        b=voluntary_guardrail("vol", "claims"),
        s=voluntary_guardrail_status("vol", "claims"),
    )
    expected = [
        scalar.voluntary_guardrail(v, 1_000.0) for v in (1_150.0, 850.0, 1_050.0)
    ]
    assert [(r["b"], r["s"]) for r in rows[:3]] == expected
    assert (rows[3]["b"], rows[3]["s"]) == (None, None)
    cats = {
        "cat": ["A&D", "High Needs", "High Needs", "High Needs"],
        "by3": ["A&D", "High Needs", "A&D", None],
    }
    flags = backend.select(
        cats,
        {"cat": pl.String, "by3": pl.String},
        f=contributes_to_voluntary_benchmark("cat", "by3"),
    )
    assert [r["f"] for r in flags] == [True, True, False, False]


def test_table_3_5_4_as_a_frame(backend):
    rows = []
    for category, (
        pbpm,
        trend,
        risk,
        full,
        hybrid,
        factor,
        _,
    ) in TABLE_3_5_4_CLAIMS.items():
        rows.append(
            {
                "category": category,
                "pbpm": pbpm,
                "trend": trend,
                "risk": risk,
                "full": float(full),
                "m2": float(hybrid),
                "f2": factor,
                "m3": 0.0,
                "f3": 1.0,
                "m4": 0.0,
                "f4": 1.0,
            }
        )
    for category, (
        pbpm,
        trend,
        risk,
        full,
        months,
        factors,
        _,
    ) in TABLE_3_5_4_VOLUNTARY.items():
        rows.append(
            {
                "category": category,
                "pbpm": pbpm,
                "trend": trend,
                "risk": risk,
                "full": float(full),
                **{f"m{q}": float(m) for q, m in zip((2, 3, 4), months, strict=True)},
                **{f"f{q}": f for q, f in zip((2, 3, 4), factors, strict=True)},
            }
        )
    py_pbpm = performance_year_benchmark_pbpm("pbpm", nw.lit(0.0), "trend", "risk")
    dollars = benchmark_dollars(
        py_pbpm, "full", [("m2", "f2"), ("m3", "f3"), ("m4", "f4")]
    )
    cells = backend.select(rows, category=nw.col("category"), dollars=dollars)
    expected = [v[-1] for v in TABLE_3_5_4_CLAIMS.values()] + [
        v[-1] for v in TABLE_3_5_4_VOLUNTARY.values()
    ]
    assert [c["dollars"] for c in cells] == expected
    line24 = {}
    for cell in cells:
        line24[cell["category"]] = scalar.total_benchmark(
            [line24.get(cell["category"], 0.0), cell["dollars"]]
        )
    assert line24 == {k: v for k, v in TABLE_3_5_4_TOTALS["24"].items() if k != "TOTAL"}


def test_benchmark_dollars_without_hybrid_months(backend):
    (result,) = backend.select(
        {"pbpm": [1_125.0], "months": [60_000.0]}, d=benchmark_dollars("pbpm", "months")
    )
    assert result["d"] == 67_500_000.0


def test_month_weighted_mean_in_a_group_by(backend):
    data = {
        "cat": ["A&D", "A&D", "A&D", "ESRD", "HN"],
        "score": [1.0, 1.2, None, 0.9, 1.5],
        "months": [10.0, 30.0, 50.0, 12.0, 0.0],
    }
    frame = nw.from_native(
        backend.frame(
            data, {"cat": pl.String, "score": pl.Float64, "months": pl.Float64}
        )
    )
    out = (
        frame.group_by("cat")
        .agg(avg=month_weighted_mean("score", "months"))
        .sort("cat")
    )
    rows = backend.rows(out.to_native())
    assert [r["cat"] for r in rows] == ["A&D", "ESRD", "HN"]
    assert rows[0]["avg"] == pytest.approx(
        scalar.month_weighted_average([(1.0, 10), (1.2, 30), (None, 50)])
    )
    assert rows[1]["avg"] == pytest.approx(0.9)
    assert rows[2]["avg"] is None