stack/src/lead/express/benchmark.py

"""Performance Year Benchmark formulas (§3) as narwhals expressions.

Row-wise builders for a frame with one row per benchmark cell (category ×
alignment type, county, base year…), plus :func:`month_weighted_mean` for
``group_by(...).agg(...)``. Scalar twins: ``lead.benchmark``. A null or
non-positive denominator gives null, never a stand-in value.

Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.
"""

from __future__ import annotations

from collections.abc import Sequence

import narwhals as nw

from lead.benchmark import (
    CAPPED,
    CAPPED_AT_UPPER,
    FLOORED,
    RAISED_TO_LOWER,
    WITHIN_CAP,
    WITHIN_GUARDRAIL,
)
from lead.express import IntoExpr, cents, col, lit_or_col, ratio
from lead.params import (
    ACPT_GUARDRAILS,
    ACPT_WEIGHT,
    FIRST_PERFORMANCE_YEAR,
    GLOBAL,
    HIGH_NEEDS,
    LOWER_SPENDING,
    REGIONAL_EFFICIENCY_SHARE,
    RISK_SCORE_GROWTH_CAPS,
    VOLUNTARY_GUARDRAIL_RATE,
)

RISK_CAP_REFERENCE_MISSING = "risk_growth_cap_reference_missing"
_NULL = nw.lit(None, dtype=nw.Float64())

# ── §3.3 ──────────────────────────────────────────────────────────────────


def risk_standardized_pbpm(
    expenditure: IntoExpr, eligible_months: IntoExpr, average_risk_score: IntoExpr
) -> nw.Expr:
    """Expenditure ÷ eligible months ÷ average risk score (Table 3.3.3
    lines 3 and 5); null without positive months and score.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.3.1 / Table 3.3.3.
    """
    return ratio(ratio(expenditure, eligible_months), average_risk_score)


def growth_rate(later: IntoExpr, earlier: IntoExpr) -> nw.Expr:
    """later ÷ earlier FFS expenditures; null unless both are positive.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) Box 3.3.2 Steps 1–2.
    """
    return nw.when(col(later) > 0).then(ratio(later, earlier))


def county_market_share(aligned_months: IntoExpr, eligible_months: IntoExpr) -> nw.Expr:
    """ACO-aligned ÷ alignment-eligible months in a county.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) Box 3.3.2 Step 3.1.
    """
    return ratio(aligned_months, eligible_months)


def county_weight(aligned_months: IntoExpr, aco_aligned_months: IntoExpr) -> nw.Expr:
    """A county's share of the ACO's aligned months (trend weights and
    regional FFS expenditures alike).

    Source: LEAD PY2027 methodology (bib E4PEU3HV) Box 3.3.2 Step 3.2 / §3.4.1.
    """
    return ratio(aligned_months, aco_aligned_months)


def two_way_blended_trend(
    national_growth: IntoExpr, regional_growth: IntoExpr, national_weight: IntoExpr
) -> nw.Expr:
    """national × weight + regional × (1 − weight).

    Source: LEAD PY2027 methodology (bib E4PEU3HV) Box 3.3.2 Step 4.
    """
    weight = col(national_weight)
    return col(national_growth) * weight + col(regional_growth) * (1.0 - weight)


def trended_weighted_pbpm(
    risk_standardized: IntoExpr, trend_to_by3: IntoExpr, base_year_weight: IntoExpr
) -> nw.Expr:
    """One base year's contribution to the Historical Benchmark: PBPM ×
    trend to BY 3 × weight (Table 3.3.3 line 9); sum it over base years.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.3.3 / Table 3.3.3.
    """
    return col(risk_standardized) * col(trend_to_by3) * col(base_year_weight)


# ── §3.2.5 ────────────────────────────────────────────────────────────────


def voluntary_guardrail(
    voluntary_historical_pbpm: IntoExpr,
    claims_historical_pbpm: IntoExpr,
    rate: float = VOLUNTARY_GUARDRAIL_RATE,
) -> nw.Expr:
    """The voluntary Historical Benchmark held within ±10 percent of the
    claims-aligned one; null when either is null.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.2.5.
    """
    claims = col(claims_historical_pbpm)
    return nw.when(~claims.is_null()).then(
        col(voluntary_historical_pbpm).clip(
            lower_bound=claims * (1.0 - rate), upper_bound=claims * (1.0 + rate)
        )
    )


def voluntary_guardrail_status(
    voluntary_historical_pbpm: IntoExpr,
    claims_historical_pbpm: IntoExpr,
    rate: float = VOLUNTARY_GUARDRAIL_RATE,
) -> nw.Expr:
    """Which side of the guardrail the voluntary benchmark landed on; null
    when either input is null.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.2.5.
    """
    voluntary, claims = col(voluntary_historical_pbpm), col(claims_historical_pbpm)
    return (
        nw.when(voluntary.is_null() | claims.is_null())
        .then(nw.lit(None, dtype=nw.String()))
        .when(voluntary > claims * (1.0 + rate))
        .then(nw.lit(CAPPED_AT_UPPER))
        .when(voluntary < claims * (1.0 - rate))
        .then(nw.lit(RAISED_TO_LOWER))
        .otherwise(nw.lit(WITHIN_GUARDRAIL))
    )


def contributes_to_voluntary_benchmark(
    category: IntoExpr, by3_category: IntoExpr
) -> nw.Expr:
    """A voluntarily aligned beneficiary contributes unless High Needs in
    the PY without being High Needs in BY 3 (a null BY 3 category is not
    High Needs).

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.2.5.
    """
    was_high_needs = (col(by3_category) == HIGH_NEEDS).fill_null(False)
    return (col(category) != HIGH_NEEDS) | was_high_needs


# ── §3.4 ──────────────────────────────────────────────────────────────────


def regional_efficiency_eligible(
    spending_designation: IntoExpr, risk_option: IntoExpr
) -> nw.Expr:
    """Lower-spending ACOs in the Global Risk track.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.4.3.
    """
    return (col(spending_designation) == LOWER_SPENDING) & (col(risk_option) == GLOBAL)


def uncapped_regional_efficiency_adjustment(
    regional_ffs_pbpm: IntoExpr, historical_benchmark_pbpm: IntoExpr
) -> nw.Expr:
    """50 percent of the positive regional − historical difference.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.4.3 / Table 3.4.3.1.
    """
    difference = col(regional_ffs_pbpm) - col(historical_benchmark_pbpm)
    return difference.clip(lower_bound=0.0) * REGIONAL_EFFICIENCY_SHARE


def benchmark_adjustment_cap(
    uspcc_amount: IntoExpr, py_average_risk_score: IntoExpr, cap_rate: IntoExpr | float
) -> nw.Expr:
    """USPCC-based amount × PY average risk score × cap rate.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.4.3 / Table 3.4.3.3.
    """
    return col(uspcc_amount) * col(py_average_risk_score) * lit_or_col(cap_rate)


def regional_efficiency_adjustment(
    regional_ffs_pbpm: IntoExpr,
    historical_benchmark_pbpm: IntoExpr,
    *,
    eligible: IntoExpr,
    uspcc_amount: IntoExpr,
    py_average_risk_score: IntoExpr,
    cap_rate: IntoExpr | float,
) -> nw.Expr:
    """Applied REA: zero when ineligible or not positive, else the lesser
    of the uncapped REA and the cap; null when a needed input is null (a
    positive REA never falls back to its uncapped amount).

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.4.3 / Tables
    3.4.3.1–3.4.3.4.
    """
    uncapped = uncapped_regional_efficiency_adjustment(
        regional_ffs_pbpm, historical_benchmark_pbpm
    )
    cap = benchmark_adjustment_cap(uspcc_amount, py_average_risk_score, cap_rate)
    flag = col(eligible)
    return (
        nw.when(uncapped.is_null() | flag.is_null())
        .then(_NULL)
        .when(~flag | (uncapped == 0))
        .then(nw.lit(0.0))
        .when(cap.is_null())
        .then(_NULL)
        .otherwise(nw.min_horizontal(uncapped, cap))
    )


# ── §3.5 ──────────────────────────────────────────────────────────────────


def _guardrail(performance_year: IntoExpr | int, side: int) -> nw.Expr:
    year = lit_or_col(performance_year)
    bound = None
    for model_year, bounds in ACPT_GUARDRAILS.items():
        condition = year == FIRST_PERFORMANCE_YEAR + model_year - 1
        bound = (nw.when if bound is None else bound.when)(condition).then(
            nw.lit(bounds[side])
        )
    return bound.otherwise(_NULL)


def three_way_blended_trend(
    two_way_trend: IntoExpr, acpt: IntoExpr, performance_year: IntoExpr | int
) -> nw.Expr:
    """Two-thirds two-way trend + one-third ACPT, held within the PY's
    guardrail around the two-way trend; null for a PY outside the model.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.5.1 / Table 3.5.1.
    """
    two_way = col(two_way_trend)
    raw = (1.0 - ACPT_WEIGHT) * two_way + ACPT_WEIGHT * col(acpt)
    lower = two_way + _guardrail(performance_year, 0)
    upper = two_way + _guardrail(performance_year, 1)
    return nw.when(~lower.is_null()).then(
        raw.clip(lower_bound=lower, upper_bound=upper)
    )


def risk_score_growth_cap(category: IntoExpr) -> nw.Expr:
    """The category's growth cap: 3 percent A&D and ESRD, 4 percent High
    Needs; null for anything else.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.5.2.
    """
    return col(category).replace_strict(
        dict(RISK_SCORE_GROWTH_CAPS), default=None, return_dtype=nw.Float64()
    )


def capped_risk_score(
    by3_risk_score: IntoExpr, py_risk_score: IntoExpr, category: IntoExpr
) -> nw.Expr:
    """PY risk score held within BY 3 × (1 ± cap). Null without a BY 3
    score or a known category: the uncapped score never passes through.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.5.2.
    """
    by3, cap = col(by3_risk_score), risk_score_growth_cap(category)
    return nw.when(~by3.is_null() & ~cap.is_null()).then(
        col(py_risk_score).clip(
            lower_bound=by3 * (1.0 - cap), upper_bound=by3 * (1.0 + cap)
        )
    )


def risk_growth_cap_status(
    by3_risk_score: IntoExpr, py_risk_score: IntoExpr, category: IntoExpr
) -> nw.Expr:
    """Whether the PY risk score was capped, raised to the floor, kept, or
    could not be capped for a missing reference.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.5.2.
    """
    by3, py, cap = (
        col(by3_risk_score),
        col(py_risk_score),
        risk_score_growth_cap(category),
    )
    return (
        nw.when(by3.is_null() | py.is_null() | cap.is_null())
        .then(nw.lit(RISK_CAP_REFERENCE_MISSING))
        .when(py > by3 * (1.0 + cap))
        .then(nw.lit(CAPPED))
        .when(py < by3 * (1.0 - cap))
        .then(nw.lit(FLOORED))
        .otherwise(nw.lit(WITHIN_CAP))
    )


def performance_year_benchmark_pbpm(
    historical_benchmark_pbpm: IntoExpr,
    benchmark_adjustment: IntoExpr,
    trend_factor: IntoExpr,
    capped_py_risk_score: IntoExpr,
) -> nw.Expr:
    """(Historical Benchmark + REA or PSA) × three-way trend × capped PY
    risk score; a null adjustment stays null.

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.5 / §3.5.2.
    """
    adjusted = col(historical_benchmark_pbpm) + col(benchmark_adjustment)
    return adjusted * col(trend_factor) * col(capped_py_risk_score)


def benchmark_dollars(
    py_benchmark_pbpm: IntoExpr,
    full_year_months: IntoExpr,
    hybrid_months_and_factors: Sequence[tuple[IntoExpr, IntoExpr]] = (),
) -> nw.Expr:
    """cents(PBPM × [full-year months + Σ hybrid months × factor]).

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.5.3 / Table 3.5.3.2 /
    Table 3.5.4 lines 8 and 22.
    """
    hybrid = nw.lit(0.0)
    for months, factor in hybrid_months_and_factors:
        hybrid = hybrid + col(months) * col(factor)
    return cents(col(py_benchmark_pbpm) * (col(full_year_months) + hybrid))


def month_weighted_mean(value: IntoExpr, months: IntoExpr) -> nw.Expr:
    """Aggregation: Σ value × months ÷ Σ months, over rows whose value is
    known and months positive; null without such rows. Use inside
    ``group_by(...).agg(...)`` (average risk scores, the claims/voluntary
    PBPM combination, regional FFS expenditure).

    Source: LEAD PY2027 methodology (bib E4PEU3HV) §3.3.1 / §3.4.1 /
    §3.5.2 / §3.5.4.
    """
    amount, weight = col(value), col(months)
    counted = ~amount.is_null() & (weight > 0)
    numerator = nw.when(counted).then(amount * weight).otherwise(nw.lit(0.0)).sum()
    denominator = nw.when(counted).then(weight).otherwise(nw.lit(0.0)).sum()
    return ratio(numerator, denominator)