stack/tests/lead/test_express_settlement.py

"""lead.express.settlement settles frames of scenarios on polars and DuckDB:
Table 5.6.1 to the cent, bit-for-bit parity with the scalar engine over a
grid of scenarios, and nulls — not zeros — for missing inputs."""

from __future__ import annotations

import random
from dataclasses import asdict

import polars as pl
import pytest

from lead.express import cents
from lead.express.settlement import (
    INPUT_COLUMNS,
    LINE_COLUMNS,
    discount_amount,
    discount_rate,
    settle,
)
from lead.money import cents as scalar_cents
from lead.params import DISCOUNT_RATES
from lead.settlement import SettlementInputs
from lead.settlement import settle as settle_scalar
from tests.lead.paper import TABLE_5_6_1, TABLE_5_6_1_INPUTS, TABLE_5_6_1_LINES

SCHEMA = {
    "py_benchmark": pl.Float64,
    "risk_option": pl.String,
    "spending_designation": pl.String,
    "performance_year": pl.Int64,
    "quality_score": pl.Float64,
    "capitation_payments": pl.Float64,
    "participant_provider_claims": pl.Float64,
    "preferred_provider_claims": pl.Float64,
    "non_aco_provider_claims": pl.Float64,
    "stop_loss_elected": pl.Boolean,
    "stop_loss_charge": pl.Float64,
    "stop_loss_payout": pl.Float64,
    "stop_loss_neutrality_factor": pl.Float64,
    "retention_incentive_applies": pl.Boolean,
}


def row(**changes):
    return {**TABLE_5_6_1_INPUTS, "risk_option": "Global", **changes}


def settled(backend, rows):
    return backend.rows(settle(backend.frame(rows, SCHEMA)))


def test_schema_matches_the_input_model():
    assert tuple(SCHEMA) == INPUT_COLUMNS
    assert "shared_savings_losses_after_sequestration" in LINE_COLUMNS


def test_table_5_6_1_every_line_on_each_backend(backend):
    rows = settled(
        backend, [row(risk_option="Global"), row(risk_option="Professional")]
    )
    for result in rows:
        expected = TABLE_5_6_1[result["risk_option"]]
        for line, field in TABLE_5_6_1_LINES:
            assert result[field] == expected[line], (result["risk_option"], line)


def scenario_grid() -> list[dict]:
    """Valid scenarios across options, designations, years, quality,
    elections, and gross results from deep losses to deep savings."""
    rng = random.Random(2027)
    scenarios = []
    for i in range(160):
        benchmark = rng.choice([150_000_000.0, 48_250_731.17, 9_999_999.99])
        option = rng.choice(["Global", "Professional"])
        elected = rng.random() < 0.7
        share = rng.choice(
            [
                -0.62,
                -0.5,
                -0.35,
                -0.2,
                -0.15,
                -0.1,
                -0.0001,
                0.0,
                0.07,
                0.1,
                0.15,
                0.2,
                0.35,
                0.5,
                0.71,
            ]
        )
        payout = rng.choice([0.0, 1_234_567.89, 2_900_000.0])
        scenarios.append(
            row(
                py_benchmark=benchmark,
                risk_option=option,
                spending_designation=rng.choice(
                    [
                        "Lower-spending",
                        "Higher-spending",
                        None if option == "Professional" else "Lower-spending",
                    ]
                ),
                performance_year=rng.choice([2027, 2028, 2030, 2031, 2036]),
                quality_score=rng.choice([0.0, 0.333, 0.95, 1.0]),
                capitation_payments=round(benchmark * rng.uniform(0.0, 0.2), 2),
                participant_provider_claims=round(benchmark * rng.uniform(0.0, 0.1), 2),
                preferred_provider_claims=round(benchmark * rng.uniform(0.0, 0.2), 2),
                non_aco_provider_claims=round(benchmark * (0.5 - share), 2),
                stop_loss_elected=elected,
                stop_loss_charge=round(benchmark * 0.02, 2) if elected else None,
                stop_loss_payout=payout if elected else None,
                stop_loss_neutrality_factor=(
                    rng.choice([0.93, 1.07, None])
                    if payout == 0.0
                    else rng.choice([0.93, 1.07])
                )
                if elected
                else None,
                retention_incentive_applies=i % 7 == 0,
            )
        )
    return scenarios


def test_frame_matches_scalar_bit_for_bit(backend):
    scenarios = scenario_grid()
    rows = settled(backend, scenarios)
    assert len(rows) == len(scenarios)
    signs = {
        (r["shared_savings_losses"] > 0) - (r["shared_savings_losses"] < 0)
        for r in rows
    }
    assert signs == {1, -1}  # savings and losses both exercised
    for scenario, result in zip(scenarios, rows, strict=True):
        expected = asdict(settle_scalar(SettlementInputs(**scenario)))
        assert {k: result[k] for k in expected} == expected


def test_other_columns_pass_through(backend):
    frame = backend.frame([{**row(), "aco": "A1"}], {**SCHEMA, "aco": pl.String})
    assert backend.rows(settle(frame))[0]["aco"] == "A1"


def test_not_elected_rows_get_the_papers_zeros(backend):
    (result,) = settled(
        backend,
        [
            row(
                stop_loss_elected=False,
                stop_loss_charge=None,
                stop_loss_payout=None,
                stop_loss_neutrality_factor=None,
            )
        ],
    )
    assert (result["stop_loss_charge"], result["stop_loss_payout"]) == (0.0, 0.0)
    assert result["stop_loss_neutrality_factor"] is None
    assert result["net_stop_loss_impact"] == 0.0
    assert result["py_expenditure_after_stop_loss"] == result["py_expenditure"]


@pytest.mark.parametrize(
    "changes,first_null",
    [
        ({"quality_score": None}, "quality_earn_back"),
        ({"quality_score": 1.2}, "quality_earn_back"),
        ({"spending_designation": None}, "discount_rate"),
        ({"spending_designation": "lower-spending"}, "discount_rate"),
        ({"performance_year": 2037}, "discount_rate"),
        ({"risk_option": "global"}, "discount_rate"),
        ({"stop_loss_charge": None}, "net_stop_loss_impact"),
        ({"stop_loss_neutrality_factor": None}, "adjusted_stop_loss_payout"),
        ({"stop_loss_elected": None}, "adjusted_stop_loss_payout"),
        ({"retention_incentive_applies": None}, "retention_incentive_adjustment"),
        ({"non_aco_provider_claims": None}, "total_ffs_payments"),
    ],
)
def test_missing_inputs_null_the_chain_never_zero_it(backend, changes, first_null):
    (result,) = settled(backend, [row(**changes)])
    assert result[first_null] is None
    assert result["shared_savings_losses_after_sequestration"] is None
    assert result["sequestration_amount"] is None


def test_missing_columns_are_named(backend):
    frame = backend.frame([{"py_benchmark": 1.0}])
    with pytest.raises(ValueError, match="quality_score"):
        settle(frame)


def test_discount_rate_lookup(backend):
    data = [
        {"o": "Global", "d": d, "y": y}
        for d, by_year in DISCOUNT_RATES.items()
        for y in by_year
    ] + [
        {"o": "Professional", "d": None, "y": 2030},
        {"o": "Professional", "d": "Lower-spending", "y": 2026},
        {"o": "Global", "d": "Lower-spending", "y": 2026},
    ]
    rates = [r["rate"] for r in backend.select(data, rate=discount_rate("o", "d", "y"))]
    assert rates[:-3] == [
        r for by_year in DISCOUNT_RATES.values() for r in by_year.values()
    ]
    assert rates[-3:] == [0.0, None, None]


def test_missing_discount_rate_is_zero_in_the_na_helper(backend):
    data = {"benchmark": [150_000_000.0, 150_000_000.0], "rate": [None, 0.03]}
    result = backend.select(
        data,
        {"benchmark": pl.Float64, "rate": pl.Float64},
        amount=discount_amount("benchmark", "rate"),
    )
    assert [r["amount"] for r in result] == [0.0, 4_500_000.0]


def test_cents_expr_matches_the_scalar_bit_for_bit(backend):
    rng = random.Random(5)
    values = [rng.uniform(-1e9, 1e9) for _ in range(500)]
    values += [k / 1000 for k in range(-2000, 2001, 5)]  # half-cent ties
    values += [0.125, -0.125, 2.675, 1.005, -0.001, 0.0, None]
    result = backend.select({"x": values}, {"x": pl.Float64}, y=cents("x"))
    for value, got in zip(values, (r["y"] for r in result), strict=True):
        assert got == (None if value is None else scalar_cents(value)), value