Causally grounded Cortex Agent that refuses when the data cannot answer: code¶

"""Causal decision console.

Runs as a Streamlit app inside Snowflake. One file, no use-case-specific code: every
label, cost, and protected attribute comes from MODELS.SPEC_SNAPSHOT, so the same
app serves any use case built on this framework.

THE ARCHITECTURAL POINT WORTH NOTICING

This app does not compute anything. It calls five of the six stored procedures in
the AGENT schema that the Cortex Agent calls, and reads the same MODELS tables the
notebooks wrote. That is deliberate and it is the whole reason the design holds
together: a figure shown here, a figure quoted in chat, and a figure printed in a
notebook are the same figure, because there is exactly one implementation of each.
The moment a console recomputes "net value" its own way, you have two definitions of
net value and no way to know which one a screenshot came from.

THE REFUSAL IS ENFORCED IN THE UI, NOT CAPTIONED

When the selected cohort cannot support an estimate, this app does not show the
numbers with a warning next to them. It withholds them. A warning beside a number is
read as a formality; an absent number is read as an answer. That behaviour is the
point of the whole pipeline, so it gets to shape the layout.
"""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path

import altair as alt
import pandas as pd
import streamlit as st


def _connect():
    """Return a session that has a current database, in either environment.

    Every query below uses schema-qualified but database-UNQUALIFIED names
    (MODELS.CATE_SCORES, DATA.UNIT, AGENT.RUN_SCENARIO, INFORMATION_SCHEMA.COLUMNS),
    which is what lets one app file serve every use case.

    In Streamlit in Snowflake that resolves for free: the session inherits the
    database of the STREAMLIT object, which deploy_app.py creates inside the use
    case's own database. Running locally there is no such context, and the first
    query fails with "This session does not have a current database" before any UI
    is built. So set it explicitly, from CAUSAL_DATABASE if given, otherwise from the
    spec on disk.
    """
    try:
        from snowflake.snowpark.context import get_active_session

        return get_active_session(), True
    except Exception:
        from snowflake.snowpark import Session

        # Same precedence as causal_core.session: env vars, otherwise the default
        # connection in connections.toml. Never a hard-coded name.
        connection = os.environ.get("SNOWFLAKE_CONNECTION_NAME") or os.environ.get(
            "SNOWFLAKE_DEFAULT_CONNECTION_NAME"
        )
        builder = Session.builder
        if connection:
            builder = builder.config("connection_name", connection)
        session = builder.getOrCreate()

        database = os.environ.get("CAUSAL_DATABASE")
        if not database:
            sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "framework"))
            from causal_core import load_spec

            database = load_spec().database
        # Snowpark validates the name rather than splicing it into SQL text.
        session.use_database(database)
        return session, False


SESSION, IN_SNOWFLAKE = _connect()

if not SESSION.get_current_database():
    raise RuntimeError(
        "This session has no current database, so the console's unqualified table "
        "references cannot resolve. In Snowflake, create the STREAMLIT object inside "
        "the use case database (framework/deploy_app.py does this). Locally, set "
        "CAUSAL_DATABASE."
    )


st.set_page_config(page_title="Causal decision console", layout="wide")


# ---------------------------------------------------------------------------
# Spec and data loading
# ---------------------------------------------------------------------------


@st.cache_data(ttl=600, show_spinner=False)
def load_spec_row() -> dict:
    """Read the spec snapshot. Everything domain-specific in this app comes from here."""
    rows = SESSION.sql(
        "SELECT * FROM MODELS.SPEC_SNAPSHOT ORDER BY SNAPSHOT_AT DESC LIMIT 1"
    ).collect()
    if not rows:
        st.error(
            "MODELS.SPEC_SNAPSHOT is empty. Run causal_core.write_spec_snapshot() or "
            "framework/deploy_app.py before opening the console."
        )
        st.stop()
    return rows[0].as_dict()


@st.cache_data(ttl=600, show_spinner=False)
def sql_df(query: str) -> pd.DataFrame:
    return SESSION.sql(query).to_pandas()


@st.cache_data(ttl=600, show_spinner=False)
def unit_columns() -> pd.DataFrame:
    return SESSION.sql(
        """
        SELECT COLUMN_NAME, DATA_TYPE
        FROM INFORMATION_SCHEMA.COLUMNS
        WHERE TABLE_SCHEMA = 'DATA' AND TABLE_NAME = 'UNIT'
        ORDER BY ORDINAL_POSITION
        """
    ).to_pandas()


# The procedures the console may call. The name is checked against this set and
# every argument is a bind parameter, so nothing a user picks in the sidebar (a
# category value such as O'Neil, for example) is ever spliced into SQL text.
PROCEDURES = {
    "CHECK_IDENTIFICATION",
    "RUN_SCENARIO",
    "OPTIMIZE_ALLOCATION",
    "COMPARE_TARGETING_STRATEGIES",
    "EXPLAIN_ASSUMPTIONS",
}


def call_proc_raw(name: str, *args):
    """Call an AGENT procedure with bound arguments and return its raw result."""
    if name not in PROCEDURES:
        raise ValueError(f"unknown procedure {name!r}")
    placeholders = ", ".join("?" for _ in args)
    return SESSION.sql(
        f"CALL AGENT.{name}({placeholders})", params=list(args)
    ).collect()[0][0]


def call_proc(name: str, *args) -> dict:
    """Call an AGENT procedure and parse its VARIANT result.

    Not cached: these are the live decision surface, and a stale scenario is worse
    than a slow one.
    """
    raw = call_proc_raw(name, *args)
    if raw is None:
        return {}
    return json.loads(raw) if isinstance(raw, str) else raw


SPEC = load_spec_row()
UNIT_PLURAL = SPEC["UNIT_LABEL_PLURAL"]
UNIT_SINGULAR = SPEC["UNIT_LABEL_SINGULAR"]
OUTCOME_LABEL = SPEC["OUTCOME_LABEL"]
DIRECTION = SPEC["OUTCOME_DIRECTION"]
COST = float(SPEC["COST_PER_UNIT"])
PROTECTED = [a for a in (SPEC["PROTECTED_ATTRIBUTES"] or "").split(",") if a]
MAX_RATIO = float(SPEC["MAX_DISPARITY_RATIO"])


def money(value) -> str:
    if value is None:
        return "n/a"
    return f"${float(value):,.0f}"


# ---------------------------------------------------------------------------
# Sidebar: cohort and capacity
# ---------------------------------------------------------------------------

st.sidebar.title(SPEC["LABEL"])
st.sidebar.caption(
    f"{UNIT_PLURAL} · {SPEC['TREATMENT_LABEL']} at {money(COST)} each · "
    f"outcome: {OUTCOME_LABEL} ({DIRECTION})"
)

st.sidebar.subheader("Cohort")
st.sidebar.caption(
    "Narrow the population to ask a subgroup question. Some subgroups cannot be "
    "answered at all, which the console will tell you rather than guess."
)


def quote_ident(name: str) -> str:
    """Double-quote an identifier taken from INFORMATION_SCHEMA."""
    return '"' + name.replace('"', '""') + '"'


cols = unit_columns()
filterable = [
    c for c in cols.COLUMN_NAME.tolist() if c.upper() != SPEC["UNIT_ID_COLUMN"].upper()
]
chosen_cols = st.sidebar.multiselect(
    "Filter on",
    filterable,
    max_selections=3,
    help="Leave empty for the whole population.",
)

cohort: dict = {}
for column in chosen_cols:
    dtype = cols.loc[cols.COLUMN_NAME == column, "DATA_TYPE"].iloc[0]
    numeric = dtype.upper() in {"NUMBER", "FLOAT", "DECIMAL", "INTEGER", "DOUBLE"}
    # column comes from the INFORMATION_SCHEMA list above, never from free text.
    distinct = sql_df(
        f"SELECT DISTINCT {quote_ident(column)} AS V FROM DATA.UNIT ORDER BY V"
    )["V"].tolist()

    if not numeric or len(distinct) <= 12:
        picked = st.sidebar.selectbox(
            f"{column} =", ["(any)"] + distinct, key=f"eq_{column}"
        )
        if picked != "(any)":
            cohort[column] = picked
    else:
        lo, hi = float(min(distinct)), float(max(distinct))
        floor = st.sidebar.slider(f"{column} at least", lo, hi, lo, key=f"min_{column}")
        if floor > lo:
            cohort[column] = {"min": floor}

cohort_json = json.dumps(cohort) if cohort else ""

st.sidebar.subheader("Capacity")
capacity = st.sidebar.number_input(
    f"{UNIT_PLURAL} to treat", min_value=1, max_value=200_000, value=8000, step=500
)
budget = st.sidebar.number_input(
    "Budget ceiling",
    min_value=0,
    max_value=500_000_000,
    value=int(capacity * COST),
    step=100_000,
    help="Used by the allocation optimiser to decide whether budget or value binds.",
)


# ---------------------------------------------------------------------------
# The identification gate, evaluated once and applied everywhere
# ---------------------------------------------------------------------------

ident = call_proc("CHECK_IDENTIFICATION", cohort_json)
ANSWERABLE = bool(ident.get("identified")) and ident.get("ok", True)

st.title("Causal decision console")

if ident.get("error"):
    st.error(ident["error"])
    st.stop()

if not ANSWERABLE:
    st.error("This cohort cannot be answered. No estimate is available.")
    st.markdown(f"**Cohort:** {ident.get('cohort', 'n/a')}")
    st.markdown(ident.get("refusal", ""))
    left, mid, right = st.columns(3)
    left.metric(f"{UNIT_PLURAL} in cohort", f"{ident.get('n', 0):,}")
    mid.metric("treated", f"{ident.get('n_treated', 0):,}")
    right.metric(
        "treated with no comparable counterpart",
        f"{float(ident.get('share_treated_without_counterpart') or 0):.1%}",
    )
    st.info(
        "The **scenario** is withheld for this cohort on purpose: showing numbers "
        "next to a refusal invites reading them anyway. The other tabs still apply, "
        "because they describe the population and the method rather than this cohort. "
        "Widen the cohort in the sidebar to get a scenario."
    )
else:
    st.success(
        f"Cohort is answerable: {ident.get('cohort')}: {ident.get('n', 0):,} "
        f"{UNIT_PLURAL}, {ident.get('n_treated', 0):,} treated."
    )


tab_decide, tab_believe, tab_target, tab_equity, tab_receipt = st.tabs(
    ["Decide", "Why believe it", "Targeting", "Equity", "Receipt"]
)


# ---------------------------------------------------------------------------
# Decide
# ---------------------------------------------------------------------------

with tab_decide:
    # Gate THIS tab only. The remaining tabs describe the population and the method
    # rather than the selected cohort, so they stay available. A blanket st.stop()
    # here would halt the whole script and leave them blank, which reads as breakage
    # rather than as a refusal.
    scenario: dict = {}
    optimum: dict = {}
    if ANSWERABLE:
        scenario = call_proc("RUN_SCENARIO", int(capacity), cohort_json, True)
        optimum = call_proc("OPTIMIZE_ALLOCATION", float(budget))

    if not ANSWERABLE or scenario.get("refused"):
        st.error(
            "No scenario is shown for this cohort, because no defensible estimate "
            "exists for it.",
            icon="\U0001f6ab",
        )
        st.markdown(scenario.get("refusal") or ident.get("refusal", ""))
        st.caption(
            "The other tabs remain available. They describe the population and the "
            "method rather than this cohort, so the refusal does not affect them."
        )
    else:
        st.subheader("If you treat the top-ranked units at this capacity")
        a, b, c, d = st.columns(4)
        a.metric(f"{UNIT_PLURAL} treated", f"{scenario.get('units_treated', 0):,}")
        b.metric("programme cost", money(scenario.get("program_cost")))
        c.metric("expected net value", money(scenario.get("expected_net_value")))
        roi = scenario.get("roi")
        d.metric("ROI", "n/a" if roi is None else f"{roi:.2f}x")

        st.caption(scenario.get("basis", ""))
        if scenario.get("note"):
            st.warning(scenario["note"])

        st.divider()
        st.subheader("How many are actually worth treating?")
        st.markdown(
            "The answer is rarely *spend the whole budget*. Marginal value falls as "
            "you move down the ranking, and the optimum is where it crosses zero."
        )
        e, f, g = st.columns(3)
        e.metric(
            "worth treating",
            f"{optimum.get('units_with_positive_expected_value', 0):,}",
        )
        f.metric(
            "recommended to treat",
            f"{optimum.get('recommended_units_to_treat', 0):,}",
        )
        g.metric("binding constraint", str(optimum.get("binding_constraint", "n/a")))
        st.info(optimum.get("interpretation", ""), icon="\U0001f4a1")

        # The shape that makes "do not spend it all" obvious at a glance.
        marginal = sql_df(
            """
            -- Cumulative net value walking down the ranking, sampled every 250 units
            -- so the chart stays small. A window function cannot be nested inside
            -- another window function's PARTITION BY, hence the CTE.
            WITH ranked AS (
                SELECT ROW_NUMBER() OVER (ORDER BY NET_VALUE_PER_UNIT DESC) AS RN,
                       SUM(NET_VALUE_PER_UNIT) OVER (
                           ORDER BY NET_VALUE_PER_UNIT DESC
                           ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
                       ) AS CUM
                FROM MODELS.CATE_SCORES
            )
            SELECT RN AS UNITS_TREATED, CUM AS CUMULATIVE_NET_VALUE
            FROM ranked
            WHERE RN % 250 = 0 OR RN = 1
            ORDER BY RN
            """
        )
        if not marginal.empty:
            peak = marginal.loc[marginal.CUMULATIVE_NET_VALUE.idxmax()]
            line = (
                alt.Chart(marginal)
                .mark_line()
                .encode(
                    x=alt.X(
                        "UNITS_TREATED:Q",
                        title=f"{UNIT_PLURAL} treated, best first",
                    ),
                    y=alt.Y(
                        "CUMULATIVE_NET_VALUE:Q",
                        title="cumulative expected net value",
                    ),
                )
            )
            at_peak = (
                alt.Chart(pd.DataFrame({"x": [peak.UNITS_TREATED]}))
                .mark_rule(strokeDash=[6, 4])
                .encode(x="x:Q")
            )
            zero = (
                alt.Chart(pd.DataFrame({"y": [0]}))
                .mark_rule(color="grey")
                .encode(y="y:Q")
            )
            st.altair_chart(line + at_peak + zero, use_container_width=True)
            st.caption(
                f"Dashed line marks the peak, at about "
                f"{int(peak.UNITS_TREATED):,} {UNIT_PLURAL}. Treating past it "
                "destroys value: those units cost more than they return."
            )


# ---------------------------------------------------------------------------
# Why believe it
# ---------------------------------------------------------------------------

with tab_believe:
    st.subheader("The same question, answered several ways")
    st.markdown(
        "These methods do not disagree because one is broken. They target different "
        "quantities under different assumptions. **The raw difference is not an effect "
        "estimate at all**; it is the selection pattern plus the effect, and it is "
        "shown only to make the naive answer visible."
    )

    diag = sql_df(
        """
        SELECT METHOD, ESTIMATE, CI_LOW, CI_HIGH, IDENTIFIED, VERDICT, N, N_TREATED
        FROM MODELS.DIAGNOSTICS ORDER BY ESTIMATE
        """
    )
    if not diag.empty:
        base = alt.Chart(diag).encode(
            y=alt.Y("METHOD:N", sort=None, title=None),
            color=alt.condition(
                alt.datum.ESTIMATE > 0
                if DIRECTION == "minimize"
                else alt.datum.ESTIMATE < 0,
                alt.value("#d62728"),  # the sign that would mean "intervention hurts"
                alt.value("#1f77b4"),
            ),
        )
        st.altair_chart(
            base.mark_rule(size=3).encode(x="CI_LOW:Q", x2="CI_HIGH:Q")
            + base.mark_point(size=110, filled=True).encode(
                x=alt.X("ESTIMATE:Q", title=f"estimated effect on {OUTCOME_LABEL}")
            )
            + alt.Chart(pd.DataFrame({"x": [0]}))
            .mark_rule(color="grey")
            .encode(x="x:Q"),
            use_container_width=True,
        )
        st.caption(
            "Red marks an estimate whose sign implies the intervention is harmful. "
            f"For a {DIRECTION} outcome, only the naive comparison should land there."
        )
        st.dataframe(diag, use_container_width=True, hide_index=True)

    st.divider()
    st.subheader("Do comparable untreated units exist?")
    st.markdown(
        "Every adjustment method depends on finding untreated units that look like "
        "treated ones. Where the two distributions do not overlap, there is nothing "
        "to compare against and no method can fix it."
    )
    overlap = sql_df(
        """
        SELECT ROUND(PROPENSITY, 2) AS PROPENSITY,
               IFF(TREATED = 1, 'treated', 'untreated') AS ARM,
               COUNT(*) AS N
        FROM MODELS.CATE_SCORES GROUP BY 1, 2 ORDER BY 1
        """
    )
    if not overlap.empty:
        st.altair_chart(
            alt.Chart(overlap)
            .mark_area(opacity=0.55, interpolate="step")
            .encode(
                x=alt.X("PROPENSITY:Q", title="estimated probability of treatment"),
                y=alt.Y("N:Q", stack=None, title=f"{UNIT_PLURAL}"),
                color=alt.Color("ARM:N", title=None),
            ),
            use_container_width=True,
        )
        st.caption(
            "Where the treated area extends beyond the untreated one, those units have "
            "no counterpart. That region is what CHECK_IDENTIFICATION refuses on."
        )

    st.divider()
    st.subheader("What the pipeline asserts about itself")
    checks = sql_df(
        "SELECT CHECK_NAME, PASSED, DETAIL FROM MODELS.CALIBRATION_CHECKS ORDER BY 1"
    )
    if not checks.empty:
        passed = int(checks.PASSED.sum())
        st.metric("calibration checks passing", f"{passed} of {len(checks)}")
        st.dataframe(checks, use_container_width=True, hide_index=True)
    st.warning(
        "Read these honestly: they verify properties of the **synthetic data**, not the "
        "correctness of the estimators. They confirm the test fixture still poses the "
        "problem it was built to pose. Estimator correctness is covered separately, "
        "by the tests in tests/."
    )

    st.divider()
    st.subheader("Did the estimate survive being attacked?")
    st.markdown(
        "Refutation tests try to break the result in ways a spurious finding would not "
        "survive. **Passing them is weak evidence for the estimate; failing any is "
        "strong evidence against it.** That asymmetry is real: refutation can find a "
        "broken result but cannot certify a sound one."
    )
    refutation = sql_df(
        """
        SELECT TEST_NAME, STATUS, BASE_ESTIMATE, NEW_ESTIMATE, RELATIVE_SHIFT, NOTE
        FROM MODELS.REFUTATION ORDER BY TEST_NAME
        """
    )
    if refutation.empty:
        st.info(
            "No refutation results yet. These run in a separate environment, because "
            "dowhy needs numpy>=2 while econml needs numpy<2. Run "
            "`causal_core.refute.refute(session, spec)`."
        )
    else:
        failed = int((refutation.STATUS != "pass").sum())
        if failed:
            st.error(f"{failed} of {len(refutation)} refutation tests FAILED.")
        else:
            st.success(
                f"All {len(refutation)} refutation tests passed; the estimate was not "
                "overturned. Weak support, not proof."
            )
        st.dataframe(refutation, use_container_width=True, hide_index=True)

    with st.expander("Full assumption statement (from the agent's own tool)"):
        st.code(call_proc_raw("EXPLAIN_ASSUMPTIONS"), language=None)


# ---------------------------------------------------------------------------
# Targeting
# ---------------------------------------------------------------------------

with tab_target:
    st.subheader("Does the ranking carry signal?")
    st.markdown(
        "Qini curves computed from **observed outcomes**, not from ground truth. This "
        "is the one evaluation here that would work on real data, which makes it the "
        "one you would actually run in production."
    )

    comparison = call_proc("COMPARE_TARGETING_STRATEGIES")
    rankings = comparison.get("rankings_evaluated", [])
    if rankings:
        cols_ = st.columns(len(rankings))
        for col, r in zip(cols_, rankings):
            col.metric(
                r["model"],
                f"{r['auuc']:+.3f}",
                r["verdict"],
                delta_color="normal" if r["auuc"] > 0 else "inverse",
            )

    qini = sql_df(
        """
        SELECT MODEL, SHARE_TARGETED, QINI, QINI_RANDOM
        FROM MODELS.QINI_CURVE ORDER BY MODEL, SHARE_TARGETED
        """
    )
    if not qini.empty:
        curves = (
            alt.Chart(qini)
            .mark_line(size=2)
            .encode(
                x=alt.X(
                    "SHARE_TARGETED:Q",
                    axis=alt.Axis(format="%"),
                    title="share of population targeted",
                ),
                y=alt.Y("QINI:Q", title=f"cumulative incremental {OUTCOME_LABEL}"),
                color=alt.Color("MODEL:N", title="ranking"),
            )
        )
        random_line = (
            alt.Chart(qini[qini.MODEL == qini.MODEL.iloc[0]])
            .mark_line(strokeDash=[6, 4], color="grey")
            .encode(x="SHARE_TARGETED:Q", y="QINI_RANDOM:Q")
        )
        st.altair_chart(curves + random_line, use_container_width=True)
        st.caption(
            "The dashed line is what a ranking with no signal traces. A curve BELOW it "
            "is worse than random, a real and common outcome for risk-score ranking, "
            "because who is worst off is a different question from who is most helped."
        )
    st.info(comparison.get("how_to_read", ""))


# ---------------------------------------------------------------------------
# Equity
# ---------------------------------------------------------------------------

with tab_equity:
    st.subheader("The fairness gate")
    st.markdown(
        "Impactability correlates with being reachable, and reachability correlates "
        "with advantage. A model optimised for measured impact will concentrate help "
        "on whoever was already easiest to help, and nothing in the accuracy metrics "
        "shows it. This is the failure mode documented in Obermeyer et al., "
        "*Science* 2019."
    )

    equity = sql_df(
        """
        SELECT POLICY, ATTRIBUTE, GROUP_VALUE, POPULATION, SELECTED, SELECTION_RATE,
               REPRESENTATION_INDEX, OBSERVED_RATIO, THRESHOLD, PASSED
        FROM MODELS.EQUITY_REPORT ORDER BY ATTRIBUTE, GROUP_VALUE
        """
    )
    if equity.empty:
        st.info("No equity report yet. Run notebook 03.")
    else:
        summary = (
            equity.groupby("ATTRIBUTE")
            .agg(
                ratio=("OBSERVED_RATIO", "max"),
                threshold=("THRESHOLD", "max"),
                passed=("PASSED", "min"),
            )
            .reset_index()
        )
        cols_ = st.columns(len(summary))
        for col, row in zip(cols_, summary.itertuples()):
            col.metric(
                row.ATTRIBUTE,
                f"{row.ratio:.2f}x",
                f"threshold {row.threshold:.2f}x: {'PASS' if row.passed else 'FAIL'}",
                delta_color="normal" if row.passed else "inverse",
            )

        attribute = st.selectbox("Break down by", sorted(equity.ATTRIBUTE.unique()))
        subset = equity[equity.ATTRIBUTE == attribute]
        st.altair_chart(
            alt.Chart(subset)
            .mark_bar()
            .encode(
                x=alt.X("GROUP_VALUE:N", title=attribute, sort=None),
                y=alt.Y(
                    "SELECTION_RATE:Q",
                    axis=alt.Axis(format="%"),
                    title="share selected for treatment",
                ),
                color=alt.condition(
                    alt.datum.REPRESENTATION_INDEX < 1,
                    alt.value("#d62728"),
                    alt.value("#1f77b4"),
                ),
            ),
            use_container_width=True,
        )
        st.caption(
            "Red bars are groups under-represented among those selected, relative to "
            "their share of the population."
        )
        st.dataframe(subset, use_container_width=True, hide_index=True)
        st.markdown(
            f"The gated policy constrains the **intersection** of "
            f"{', '.join(PROTECTED)} so every marginal disparity ratio stays under "
            f"{MAX_RATIO:.2f}x. Notebook 03 prices what that constraint costs rather "
            "than presenting fairness as free."
        )


# ---------------------------------------------------------------------------
# Receipt
# ---------------------------------------------------------------------------

with tab_receipt:
    st.subheader("What produced these numbers")
    st.markdown(
        "Everything on the other tabs comes from tables written by the notebooks, read "
        "through the same procedures the Cortex Agent uses. Nothing is recomputed here, "
        "so a screenshot from this console and an answer from the agent cannot disagree."
    )

    meta = {k: v for k, v in SPEC.items() if k != "SPEC_JSON"}
    st.dataframe(
        pd.DataFrame({"field": list(meta), "value": [str(v) for v in meta.values()]}),
        use_container_width=True,
        hide_index=True,
    )

    st.subheader("Table freshness")
    freshness = sql_df(
        """
        SELECT TABLE_NAME, ROW_COUNT, LAST_ALTERED
        FROM INFORMATION_SCHEMA.TABLES
        WHERE TABLE_SCHEMA = 'MODELS' ORDER BY LAST_ALTERED DESC
        """
    )
    st.dataframe(freshness, use_container_width=True, hide_index=True)

    st.subheader("What would survive on real data")
    st.markdown(
        """
| Capability | Needs planted truth? |
|---|---|
| Propensity, overlap, identification gate | No |
| Doubly-robust ATE and its confidence interval | No |
| Difference-in-differences and the pre-trend test | No |
| Per-unit effects and the resulting ranking | No |
| **Qini curves and AUUC** | **No**, computed from observed outcomes |
| Allocation, capacity constraints, equity gate | No |
| Projected value of a policy | No, uses estimated effects |
| *Realised* value of a policy | Yes |
| Estimator error, bias, CI coverage | Yes |

The synthetic data is a **test fixture**, not evidence. Its magnitudes are
consequences of constants chosen in the generator and should not be quoted as
findings. What it establishes is that this pipeline recovers a known answer and
refuses when it cannot.
"""
    )

    with st.expander("Full spec JSON"):
        st.json(json.loads(SPEC["SPEC_JSON"]))

"""Causal decision framework: domain-blind causal inference on Snowflake.

Nothing in this package contains a domain-specific column name, label, or
dollar figure. Domain knowledge arrives through a use case spec.

Typical notebook opening:

    from causal_core import load_spec, get_session, load_analysis_frame

    spec = load_spec()                      # reads config.yaml active_usecase
    session = get_session(spec)
    frame = load_analysis_frame(session, spec)
"""

from .allocate import (
    allocate_by_group,
    allocate_positive_value,
    allocate_top_n,
    auuc,
    benefit_scores,
    compare_policies,
    net_value_per_unit,
    policy_economics,
    qini_curve,
    write_allocation,
    write_qini,
)
from .data import (
    create_analysis_view,
    COL_OUTCOME,
    COL_PRE_OUTCOME,
    COL_TREATMENT,
    COL_UNIT_ID,
    DataError,
    analysis_frame_sql,
    design_matrix,
    feature_columns,
    load_analysis_frame,
    load_capacity,
    load_did_frame,
    load_event_study,
    load_panel,
    load_truth,
    validate_tables,
)
from .equity import (
    DisparityReport,
    EquityGateFailure,
    allocate_with_floor,
    min_achievable_disparity_ratio,
    check_equity,
    disparity_report,
    write_equity_report,
)
from .session import describe_session, get_session, resolve_connection_name
from .spec import (
    write_spec_snapshot,
    CANONICAL_TABLES,
    SCHEMAS,
    Outcome,
    SpecError,
    UseCaseSpec,
    available_usecases,
    load_config,
    load_copy,
    load_spec,
)

__all__ = [
    # spec
    "UseCaseSpec",
    "Outcome",
    "SpecError",
    "load_spec",
    "load_config",
    "load_copy",
    "available_usecases",
    "CANONICAL_TABLES",
    "SCHEMAS",
    # session
    "get_session",
    "describe_session",
    "resolve_connection_name",
    # data
    "validate_tables",
    "load_analysis_frame",
    "analysis_frame_sql",
    "load_panel",
    "load_did_frame",
    "load_event_study",
    "load_capacity",
    "load_truth",
    "feature_columns",
    "design_matrix",
    "write_spec_snapshot",
    "create_analysis_view",
    "refute",
    "refutation_table",
    "RefutationUnavailable",
    "DataError",
    # allocation
    "benefit_scores",
    "net_value_per_unit",
    "allocate_top_n",
    "allocate_by_group",
    "allocate_positive_value",
    "policy_economics",
    "compare_policies",
    "qini_curve",
    "auuc",
    "write_allocation",
    "write_qini",
    # equity
    "check_equity",
    "disparity_report",
    "allocate_with_floor",
    "min_achievable_disparity_ratio",
    "write_equity_report",
    "DisparityReport",
    "EquityGateFailure",
    "COL_UNIT_ID",
    "COL_TREATMENT",
    "COL_OUTCOME",
    "COL_PRE_OUTCOME",
]

"""Turn a ranking into a decision, and price it.

A ranked list is not a decision. The decision is which units to treat given a
fixed budget, and whether doing so is worth more than it costs. That is where an
uplift model either earns its keep or does not.

Three things live here:

**Allocation.** Selecting units under a capacity constraint, optionally per group,
optionally requiring positive net value per unit. Real-data usable: it consumes
estimated effects.

**Policy comparison.** Pricing a selection using an effect estimate. Pass the
estimated CATE to get what you would have projected; pass planted truth to get what
actually would have happened. The gap between those two is worth looking at.

**Qini curves.** The standard uplift evaluation, computed from *observed outcomes*
rather than from ground truth, so it works on real data. This is the honest way to
check whether a ranking has any signal when there is no answer key.

Nothing here knows what the outcome means. Direction and monetary value come from
the spec, so a use case that maximizes revenue and one that minimizes readmission
run through identical code.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from .data import COL_OUTCOME, COL_TREATMENT, COL_UNIT_ID
from .spec import Outcome, UseCaseSpec


# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------


def benefit_scores(
    spec: UseCaseSpec,
    effect: pd.Series,
    outcome: Outcome | None = None,
) -> pd.Series:
    """Re-express per-unit effects so that larger is always better.

    This single call is what lets everything below ignore whether the outcome is
    being minimized or maximized.
    """
    outcome = outcome or spec.outcome
    return effect.apply(outcome.benefit).rename("benefit")


def net_value_per_unit(
    spec: UseCaseSpec,
    effect: pd.Series,
    outcome: Outcome | None = None,
) -> pd.Series:
    """Expected value of treating each unit, minus what treating it costs.

    Units with a negative value here are ones you would lose money on even if the
    effect estimate is exactly right, so treating them is not a close call.
    """
    outcome = outcome or spec.outcome
    gross = benefit_scores(spec, effect, outcome) * outcome.value_per_event
    return (gross - spec.treatment.cost_per_unit).rename("net_value")


# ---------------------------------------------------------------------------
# Allocation
# ---------------------------------------------------------------------------


def allocate_top_n(scores: pd.Series, capacity: int) -> pd.Series:
    """Take the highest-scoring units up to capacity."""
    selected = pd.Series(False, index=scores.index, name="selected")
    if capacity > 0:
        selected.loc[scores.nlargest(min(capacity, len(scores))).index] = True
    return selected


def allocate_by_group(
    frame: pd.DataFrame,
    spec: UseCaseSpec,
    scores: pd.Series,
    group_capacity: dict[str, int],
) -> pd.Series:
    """Take the highest-scoring units within each group, up to that group's slots.

    This is what a real constraint usually looks like: nurses are hired per region
    and offer budgets are set per market, so capacity does not pool globally. It
    matters, because a globally optimal list is frequently unimplementable.
    """
    group_col = _resolve(frame, spec.capacity.group_column)
    selected = pd.Series(False, index=frame.index, name="selected")

    for group, slots in group_capacity.items():
        rows = frame.index[frame[group_col] == group]
        if len(rows) == 0 or slots <= 0:
            continue
        take = scores.loc[rows].nlargest(min(int(slots), len(rows))).index
        selected.loc[take] = True
    return selected


def allocate_positive_value(
    spec: UseCaseSpec,
    effect: pd.Series,
    capacity: int | None = None,
    outcome: Outcome | None = None,
) -> pd.Series:
    """Treat only units expected to return more than they cost, up to capacity.

    Distinct from top-N in a way that matters: if fewer units clear the bar than
    there is capacity for, this leaves capacity unused rather than spending it on
    units it expects to lose money on. Spending the whole budget is not the goal.
    """
    value = net_value_per_unit(spec, effect, outcome)
    worth_it = value > 0
    if capacity is None:
        return worth_it.rename("selected")
    ranked = value.where(worth_it, other=-np.inf)
    return allocate_top_n(ranked, min(capacity, int(worth_it.sum())))


# ---------------------------------------------------------------------------
# Pricing a policy
# ---------------------------------------------------------------------------


def policy_economics(
    spec: UseCaseSpec,
    selected: pd.Series,
    effect: pd.Series,
    outcome: Outcome | None = None,
    label: str = "policy",
) -> dict[str, float | str]:
    """Price a selection using a given set of per-unit effects.

    ``effect`` decides what question is being answered. Estimated effects give the
    business case you would have written. Planted truth gives what would actually
    have happened. Reporting only the first is how overconfident programs get
    funded.
    """
    outcome = outcome or spec.outcome
    chosen = selected.astype(bool)
    n = int(chosen.sum())

    if n == 0:
        return {
            "policy": label,
            "units_treated": 0,
            "avg_benefit_per_unit": 0.0,
            "total_benefit": 0.0,
            "program_cost": 0.0,
            "gross_value": 0.0,
            "net_value": 0.0,
            "roi": float("nan"),
            "cost_per_event": float("inf"),
        }

    benefit = benefit_scores(spec, effect, outcome).loc[chosen[chosen].index]
    total_benefit = float(benefit.sum())
    cost = float(n * spec.treatment.cost_per_unit)

    out: dict[str, float | str] = {
        "policy": label,
        "units_treated": n,
        "avg_benefit_per_unit": float(benefit.mean()),
        "total_benefit": total_benefit,
        "program_cost": cost,
    }
    if outcome.value_per_event is not None:
        gross = total_benefit * outcome.value_per_event
        out.update(
            gross_value=float(gross),
            net_value=float(gross - cost),
            roi=float((gross - cost) / cost) if cost else float("nan"),
            cost_per_event=float(cost / total_benefit)
            if total_benefit
            else float("inf"),
        )
    return out


def compare_policies(
    spec: UseCaseSpec,
    policies: dict[str, pd.Series],
    effect: pd.Series,
    outcome: Outcome | None = None,
) -> pd.DataFrame:
    """Price several selections side by side against the same effects."""
    rows = [
        policy_economics(spec, sel, effect, outcome, label)
        for label, sel in policies.items()
    ]
    frame = pd.DataFrame(rows)
    if "net_value" in frame.columns:
        frame = frame.sort_values("net_value", ascending=False).reset_index(drop=True)
    return frame


# ---------------------------------------------------------------------------
# Qini: evaluating a ranking without an answer key
# ---------------------------------------------------------------------------


def qini_curve(
    frame: pd.DataFrame,
    spec: UseCaseSpec,
    scores: pd.Series,
    n_points: int = 25,
    outcome: Outcome | None = None,
) -> pd.DataFrame:
    """Qini curve computed from observed outcomes, not from ground truth.

    At each depth k, take the top k units by score and compute

        qini(k) = Y_treated(k) - Y_control(k) * N_treated(k) / N_control(k)

    which is the incremental outcome attributable to treating that group, scaling
    the control arm up to the treated arm's size. Signs are applied through the
    spec so a minimized outcome reads the same way as a maximized one.

    This matters because it is the only ranking evaluation here that works on real
    data. The policy comparisons against planted truth are synthetic scaffolding; this
    is what you would actually run in production to check a model still ranks.

    Caveat worth stating out loud: this is unbiased for a randomized treatment. On
    observationally assigned treatment, as in both bundled use cases, it inherits
    the same confounding as any raw comparison. Use it to compare rankings against
    each other, not to estimate an absolute effect.
    """
    outcome = outcome or spec.outcome
    work = pd.DataFrame(
        {
            "score": scores.to_numpy(),
            "t": frame[COL_TREATMENT].to_numpy(),
            "y": frame[COL_OUTCOME].to_numpy(dtype=float) * outcome.sign,
        }
    ).sort_values("score", ascending=False, kind="mergesort")

    cum_t = work["t"].cumsum().to_numpy()
    cum_c = (1 - work["t"]).cumsum().to_numpy()
    cum_yt = (work["y"] * work["t"]).cumsum().to_numpy()
    cum_yc = (work["y"] * (1 - work["t"])).cumsum().to_numpy()

    n = len(work)
    depths = np.unique(np.linspace(1, n, n_points).astype(int))

    with np.errstate(divide="ignore", invalid="ignore"):
        qini = cum_yt - np.where(cum_c > 0, cum_yc * cum_t / np.maximum(cum_c, 1), 0.0)

    total = float(qini[-1])
    rows = []
    for k in depths:
        i = k - 1
        rows.append(
            {
                "depth": int(k),
                "share_targeted": k / n,
                "n_treated_in_group": int(cum_t[i]),
                "n_control_in_group": int(cum_c[i]),
                "qini": float(qini[i]),
                # A ranking with no signal traces this straight line.
                "qini_random": total * (k / n),
            }
        )
    return pd.DataFrame(rows)


def auuc(curve: pd.DataFrame) -> float:
    """Area between the Qini curve and the random line, normalized.

    Positive means the ranking beats random. Zero means it carries no signal.
    Negative means it is actively worse than random, which does happen and is worth
    knowing before shipping.
    """
    x = curve["share_targeted"].to_numpy()
    lift = curve["qini"].to_numpy() - curve["qini_random"].to_numpy()
    area = float(np.trapz(lift, x))
    scale = abs(float(curve["qini"].iloc[-1]))
    return area / scale if scale else 0.0


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------


def write_allocation(
    session,
    spec: UseCaseSpec,
    frame: pd.DataFrame,
    selected: pd.Series,
    effect: pd.Series,
    policy_label: str,
    table: str = "ALLOCATION",
) -> str:
    """Persist a selection so the agent and app read one shared decision.

    Written as a table rather than recomputed downstream, so the number quoted in a
    chat answer and the number on a dashboard cannot drift apart.
    """
    out = pd.DataFrame(
        {
            spec.unit.id_column: frame[COL_UNIT_ID].to_numpy(),
            "POLICY": policy_label,
            "SELECTED": selected.astype(bool).to_numpy(),
            "CATE": effect.to_numpy(dtype=float),
            "BENEFIT": benefit_scores(spec, effect).to_numpy(dtype=float),
            "NET_VALUE": net_value_per_unit(spec, effect).to_numpy(dtype=float)
            if spec.outcome.value_per_event is not None
            else np.nan,
        }
    )
    session.write_pandas(
        out,
        table,
        database=spec.database,
        schema="MODELS",
        auto_create_table=True,
        overwrite=True,
        quote_identifiers=False,
    )
    return spec.models(table)


def write_qini(
    session,
    spec: UseCaseSpec,
    curve: pd.DataFrame,
    label: str | None = None,
    table: str = "QINI_CURVE",
) -> str:
    """Persist one or more Qini curves.

    ``label`` names a single-model frame. Pass None when the frame already carries a
    MODEL column, which is how several rankings are stored together. Storing every
    ranking that was evaluated matters for the agent: asked whether risk targeting
    is good enough, it can only answer with a number if the risk curve is actually
    in the table. With only the winning curve persisted it hedges, which reads as
    evasion even though the underlying result is decisive.
    """
    out = curve.copy()
    out.columns = [c.upper() for c in out.columns]
    if label is not None:
        out["MODEL"] = label
    elif "MODEL" not in out.columns:
        raise ValueError("pass a label, or include a MODEL column in the frame")
    session.write_pandas(
        out,
        table,
        database=spec.database,
        schema="MODELS",
        auto_create_table=True,
        overwrite=True,
        quote_identifiers=False,
    )
    return spec.models(table)


def _resolve(frame: pd.DataFrame, column: str) -> str:
    """Case-insensitive column lookup, so specs need not match Snowflake casing."""
    available = {c.upper(): c for c in frame.columns}
    if column.upper() not in available:
        raise KeyError(
            f"column '{column}' is not in the analysis frame. Available: {sorted(frame.columns)}"
        )
    return available[column.upper()]

"""Layer B: estimators that adjust for selection and can be attacked.

Three methods, chosen so they lean on *different* assumptions:

``dr_ate``
    Augmented inverse propensity weighting with cross-fitting. Consistent if
    either the propensity model or the outcome model is right, which is why it is
    the headline number rather than plain regression adjustment.

``causal_forest_cate``
    econml's ``CausalForestDML``. Produces per-unit effects with honest
    intervals, which is what the allocation layer needs. This is the only
    estimator here whose output can rank units.

``did_estimate``
    Difference-in-differences on the panel. Assumes parallel pre-trends rather
    than no unobserved confounding, so it is a genuinely independent check.
    Targets the effect among the treated, not the population average.

Agreement between methods that rest on different assumptions is far stronger
evidence than a small p-value from any one of them.

Everything here uses econml and scikit-learn from the Snowflake Anaconda channel.
dowhy is NOT importable alongside econml (numpy>2 vs numpy<2); refutation lives
in :mod:`causal_core.refute`, which calls an isolated stored procedure.
"""

from __future__ import annotations

import json

import numpy as np
import pandas as pd
from sklearn.ensemble import (
    HistGradientBoostingClassifier,
    HistGradientBoostingRegressor,
)
from sklearn.model_selection import StratifiedKFold

from .data import COL_OUTCOME, COL_TREATMENT, COL_UNIT_ID, design_matrix
from .results import EstimateResult, normal_ci
from .spec import Outcome, UseCaseSpec

DEFAULT_TRIM = 0.02
DEFAULT_FOLDS = 5


# ---------------------------------------------------------------------------
# Propensity and overlap
# ---------------------------------------------------------------------------


def fit_propensity(
    frame: pd.DataFrame,
    spec: UseCaseSpec,
    n_folds: int = DEFAULT_FOLDS,
    random_state: int = 11,
) -> pd.Series:
    """Cross-fitted probability of treatment for every unit.

    Cross-fitting matters. An in-sample propensity score is overfitted toward the
    observed treatment, which pushes weights toward 1 and quietly reintroduces
    exactly the bias the weighting was supposed to remove.
    """
    X = design_matrix(spec, frame)
    t = frame[COL_TREATMENT].to_numpy()
    out = np.zeros(len(frame), dtype=float)

    splitter = StratifiedKFold(
        n_splits=n_folds, shuffle=True, random_state=random_state
    )
    for train_idx, test_idx in splitter.split(X, t):
        model = HistGradientBoostingClassifier(
            max_iter=200, learning_rate=0.08, max_depth=6, random_state=random_state
        )
        model.fit(X.iloc[train_idx], t[train_idx])
        out[test_idx] = model.predict_proba(X.iloc[test_idx])[:, 1]

    return pd.Series(out, index=frame.index, name="propensity")


def overlap_report(
    frame: pd.DataFrame,
    propensity: pd.Series,
    trim: float = DEFAULT_TRIM,
) -> dict[str, float | int]:
    """Whether comparable untreated units actually exist.

    This is the diagnostic that decides whether a question is answerable at all.
    When a cohort has almost no untreated units, no method recovers the effect and
    the honest response is to refuse rather than to report a number.
    """
    t = frame[COL_TREATMENT].to_numpy()
    e = propensity.to_numpy()

    treated_e, control_e = e[t == 1], e[t == 0]
    outside = (e < trim) | (e > 1 - trim)

    # Common support: the propensity range where both arms are represented.
    lo = max(treated_e.min(), control_e.min())
    hi = min(treated_e.max(), control_e.max())

    return {
        "n": int(len(frame)),
        "n_treated": int(t.sum()),
        "treated_share": float(t.mean()),
        "propensity_min": float(e.min()),
        "propensity_max": float(e.max()),
        "treated_propensity_mean": float(treated_e.mean()),
        "control_propensity_mean": float(control_e.mean()),
        "common_support_low": float(lo),
        "common_support_high": float(hi),
        "n_outside_trim": int(outside.sum()),
        "share_outside_trim": float(outside.mean()),
        # The single number to look at: how many treated units have no plausible
        # untreated counterpart.
        "share_treated_above_0_95": float((treated_e > 0.95).mean()),
        "trim": float(trim),
    }


def is_identified(
    overlap: dict[str, float | int], max_extreme_share: float = 0.25
) -> tuple[bool, str]:
    """Decide whether an effect is estimable, and say why if not.

    Used by the agent tools so refusal is a property of the data rather than a
    scripted behavior.
    """
    share = float(overlap["share_treated_above_0_95"])
    if share > max_extreme_share:
        return False, (
            f"{share:.1%} of treated units have a propensity above 0.95, meaning "
            "there is effectively no comparable untreated group. No adjustment "
            "method can recover the effect here, so any point estimate would be "
            "an extrapolation from the model rather than a measurement."
        )
    if int(overlap["n_treated"]) < 100:
        return False, (
            f"only {overlap['n_treated']} treated units in this cohort, which is too "
            "few for a stable estimate."
        )
    return (
        True,
        "overlap is adequate: comparable untreated units exist across the propensity range.",
    )


def balance_table(
    frame: pd.DataFrame,
    spec: UseCaseSpec,
    propensity: pd.Series,
    trim: float = DEFAULT_TRIM,
) -> pd.DataFrame:
    """Standardized mean differences before and after IPW.

    A standardized difference above about 0.1 after weighting is the conventional
    signal that a covariate is still imbalanced and the adjustment has not done
    its job.
    """
    X = design_matrix(spec, frame)
    t = frame[COL_TREATMENT].to_numpy()
    e = propensity.to_numpy().clip(trim, 1 - trim)
    w = np.where(t == 1, 1 / e, 1 / (1 - e))

    rows = []
    for col in X.columns:
        x = X[col].to_numpy()
        m1, m0 = x[t == 1].mean(), x[t == 0].mean()
        pooled = np.sqrt((x[t == 1].var() + x[t == 0].var()) / 2)
        raw = (m1 - m0) / pooled if pooled > 0 else 0.0

        w1, w0 = w[t == 1], w[t == 0]
        wm1 = np.average(x[t == 1], weights=w1)
        wm0 = np.average(x[t == 0], weights=w0)
        weighted = (wm1 - wm0) / pooled if pooled > 0 else 0.0

        rows.append(
            {
                "covariate": col,
                "smd_unadjusted": raw,
                "smd_weighted": weighted,
                "still_imbalanced": abs(weighted) > 0.10,
            }
        )
    return pd.DataFrame(rows).sort_values("smd_weighted", key=np.abs, ascending=False)


# ---------------------------------------------------------------------------
# Doubly robust ATE
# ---------------------------------------------------------------------------


def dr_ate(
    frame: pd.DataFrame,
    spec: UseCaseSpec,
    outcome: Outcome | None = None,
    propensity: pd.Series | None = None,
    trim: float = DEFAULT_TRIM,
    n_folds: int = DEFAULT_FOLDS,
    random_state: int = 11,
) -> EstimateResult:
    """Cross-fitted AIPW estimate of the average treatment effect.

    The influence-function form gives a standard error that reflects both the
    outcome model and the weighting, so unlike the naive intervals in
    :mod:`causal_core.naive` this one widens when the data is genuinely
    uninformative.
    """
    outcome = outcome or spec.outcome
    X = design_matrix(spec, frame)
    t = frame[COL_TREATMENT].to_numpy()
    y = frame[COL_OUTCOME].to_numpy(dtype=float)

    if propensity is None:
        propensity = fit_propensity(
            frame, spec, n_folds=n_folds, random_state=random_state
        )
    e = propensity.to_numpy().clip(trim, 1 - trim)

    mu1, mu0 = _cross_fit_outcome_models(X, t, y, outcome, n_folds, random_state)

    keep = (propensity.to_numpy() >= trim) & (propensity.to_numpy() <= 1 - trim)
    n_dropped = int((~keep).sum())

    # AIPW influence function.
    psi = (mu1 - mu0) + t * (y - mu1) / e - (1 - t) * (y - mu0) / (1 - e)
    psi = psi[keep]

    estimate = float(psi.mean())
    se = float(psi.std(ddof=1) / np.sqrt(len(psi)))
    low, high = normal_ci(estimate, se)

    overlap = overlap_report(frame, propensity, trim=trim)
    identified, reason = is_identified(overlap)

    notes = [
        "Doubly robust: consistent if either the propensity model or the outcome "
        "model is correctly specified, rather than requiring both.",
        "Nuisance models are cross-fitted, so the propensity scores are not "
        "overfitted toward the observed treatment.",
        f"Trimmed {n_dropped} units with propensity outside [{trim}, {1 - trim}].",
        "Still assumes no unobserved confounding. Run the refutation suite before "
        "acting on it.",
    ]
    if not identified:
        notes.append(f"OVERLAP WARNING: {reason}")

    return EstimateResult(
        method="doubly robust (AIPW)",
        outcome_column=outcome.column,
        estimate=estimate,
        se=se,
        ci_low=low,
        ci_high=high,
        baseline=float(mu0[keep].mean()),
        n=int(keep.sum()),
        n_treated=int(t[keep].sum()),
        identified=identified,
        diagnostics={
            "target": "ate",
            "n_dropped_by_trim": n_dropped,
            "identification_reason": reason,
            **overlap,
        },
        notes=notes,
    )


# ---------------------------------------------------------------------------
# Heterogeneous effects
# ---------------------------------------------------------------------------


def causal_forest_cate(
    frame: pd.DataFrame,
    spec: UseCaseSpec,
    outcome: Outcome | None = None,
    n_estimators: int = 300,
    min_samples_leaf: int = 30,
    random_state: int = 11,
) -> EstimateResult:
    """Per-unit treatment effects via econml's CausalForestDML.

    This is the estimator the decision actually depends on. The ATE can be right
    while the per-unit effects are useless, so what matters here is whether the
    per-unit estimates can rank units, which the notebook tests against a policy
    comparison rather than assuming.
    """
    from econml.dml import CausalForestDML

    outcome = outcome or spec.outcome
    X = design_matrix(spec, frame)
    t = frame[COL_TREATMENT].to_numpy()
    y = frame[COL_OUTCOME].to_numpy(dtype=float)

    # model_y must be a REGRESSOR even when the outcome is binary. DML
    # residualizes the outcome rather than classifying it, so it needs E[Y|X] as a
    # continuous prediction. Passing a classifier here fails with
    # "Cannot use a classifier as a first stage model when the target is
    # continuous!". Only the treatment model is a classifier, paired with
    # discrete_treatment=True.
    est = CausalForestDML(
        model_y=HistGradientBoostingRegressor(max_iter=150, random_state=random_state),
        model_t=HistGradientBoostingClassifier(max_iter=150, random_state=random_state),
        discrete_treatment=True,
        n_estimators=n_estimators,
        min_samples_leaf=min_samples_leaf,
        random_state=random_state,
    )
    est.fit(y, t, X=X)

    cate = pd.Series(est.effect(X).ravel(), index=frame.index, name="cate")
    lb, ub = est.effect_interval(X, alpha=0.05)
    ate = float(est.ate(X))
    ate_lb, ate_ub = est.ate_interval(X, alpha=0.05)

    return EstimateResult(
        method="causal forest (DML)",
        outcome_column=outcome.column,
        estimate=ate,
        se=None,
        ci_low=float(ate_lb),
        ci_high=float(ate_ub),
        baseline=float(frame.loc[frame[COL_TREATMENT] == 0, COL_OUTCOME].mean()),
        n=len(frame),
        n_treated=int(t.sum()),
        cate=cate,
        diagnostics={
            "target": "ate",
            "n_estimators": n_estimators,
            "cate_p10": float(np.percentile(cate, 10)),
            "cate_p50": float(np.percentile(cate, 50)),
            "cate_p90": float(np.percentile(cate, 90)),
            "cate_interval_mean_width": float(np.mean(ub - lb)),
            # Share of units whose individual interval excludes zero. If this is
            # near zero the model found no unit it can confidently say anything
            # about, regardless of how much spread the point estimates show.
            "share_individually_significant": float(np.mean((lb > 0) | (ub < 0))),
        },
        notes=[
            "Orthogonalized: outcome and treatment are both residualized before "
            "the effect is fit, so confounding enters through the nuisance models "
            "rather than the effect estimate.",
            "Per-unit intervals are honest but wide. Use the ranking, not the "
            "individual numbers, for allocation.",
        ],
    )


# ---------------------------------------------------------------------------
# Difference in differences
# ---------------------------------------------------------------------------


def did_estimate(
    did_frame: pd.DataFrame,
    spec: UseCaseSpec,
    event_study: pd.DataFrame | None = None,
) -> EstimateResult:
    """Difference-in-differences from unit-level pre/post means.

    Takes the frame produced by :func:`causal_core.data.load_did_frame`, which has
    one row per unit with its mean outcome before and after treatment.

    Comparing the change in those means across arms is algebraically identical to
    a two-way fixed effects regression on the full panel, but it runs on
    thousands of rows instead of millions, and the standard error is already
    clustered at the unit level because each unit contributes one observation.

    Rests on parallel pre-trends rather than on no unobserved confounding, so it
    is a genuinely independent check on the AIPW estimate rather than a
    restatement of it. Targets the effect among the treated.
    """
    treated = did_frame.loc[did_frame[COL_TREATMENT] == 1, "DELTA"].dropna()
    control = did_frame.loc[did_frame[COL_TREATMENT] == 0, "DELTA"].dropna()

    estimate = float(treated.mean() - control.mean())
    se = float(
        np.sqrt(treated.var(ddof=1) / len(treated) + control.var(ddof=1) / len(control))
    )
    low, high = normal_ci(estimate, se)

    notes = [
        "Assumes parallel pre-treatment trends. That is a different assumption "
        "from no unobserved confounding, so agreement with AIPW is meaningful.",
        "Unit-level differences, so the standard error is clustered by unit by "
        "construction rather than by a correction.",
    ]
    diagnostics: dict[str, float | str | int] = {
        "target": "att",
        "n_treated_units": int(len(treated)),
        "n_control_units": int(len(control)),
        "pre_mean_treated": float(
            did_frame.loc[did_frame[COL_TREATMENT] == 1, "PRE_MEAN"].mean()
        ),
        "pre_mean_control": float(
            did_frame.loc[did_frame[COL_TREATMENT] == 0, "PRE_MEAN"].mean()
        ),
    }

    if event_study is not None and len(event_study):
        pre_trend = _pre_trend_test(event_study)
        diagnostics.update(
            pre_trend_slope=pre_trend["slope"], pre_trend_p_value=pre_trend["p_value"]
        )
        notes.append(
            f"Pre-trend test: the treated-minus-control gap moves by "
            f"{pre_trend['slope']:+.5f} per period before treatment "
            f"(p={pre_trend['p_value']:.3f})."
        )
        if pre_trend["p_value"] < 0.05:
            notes.append(
                "PRE-TREND WARNING: the gap was already moving before treatment, "
                "which undermines the parallel trends assumption."
            )

    return EstimateResult(
        method="difference in differences",
        outcome_column=spec.outcome.column,
        estimate=estimate,
        se=se,
        ci_low=low,
        ci_high=high,
        baseline=float(
            did_frame.loc[did_frame[COL_TREATMENT] == 0, "POST_MEAN"].mean()
        ),
        n=int(len(did_frame)),
        n_treated=int(len(treated)),
        diagnostics=diagnostics,
        notes=notes,
    )


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------


def write_cate_scores(
    session,
    spec: UseCaseSpec,
    frame: pd.DataFrame,
    result: EstimateResult,
    table: str = "CATE_SCORES",
    propensity: pd.Series | None = None,
) -> str:
    """Persist per-unit effects so the agent and app never refit a model.

    This is the handoff between the econml world and everything else. Data moves
    through tables rather than pickles, which is mandatory here because the dowhy
    procedure runs on a different numpy and pandas major version.

    ``propensity`` is written alongside because it is what makes cohort-level
    refusal possible downstream. An agent asked about a narrow slice needs to know
    whether that slice has common support, and that question cannot be answered
    from an average treatment effect or a per-unit CATE alone. Without this column
    the agent can only report the whole-population verdict, which is exactly the
    kind of overconfident generalisation the gate exists to prevent.
    """
    if result.cate is None:
        raise ValueError(f"{result.method} produced no per-unit effects to write")

    out = pd.DataFrame(
        {
            spec.unit.id_column: frame[COL_UNIT_ID].to_numpy(),
            "CATE": result.cate.to_numpy(dtype=float),
            "BENEFIT": [
                spec.outcome.benefit(v) for v in result.cate.to_numpy(dtype=float)
            ],
            "METHOD": result.method,
            "OUTCOME_COLUMN": result.outcome_column,
        }
    )
    if propensity is not None:
        out["PROPENSITY"] = propensity.to_numpy(dtype=float)
        out["TREATED"] = frame[COL_TREATMENT].to_numpy(dtype=int)
    out["VALUE_PER_UNIT"] = (
        out["BENEFIT"] * spec.outcome.value_per_event
        if spec.outcome.value_per_event is not None
        else np.nan
    )
    out["COST_PER_UNIT"] = spec.treatment.cost_per_unit
    out["NET_VALUE_PER_UNIT"] = out["VALUE_PER_UNIT"] - out["COST_PER_UNIT"]

    fqn = spec.models(table)
    session.write_pandas(
        out,
        table,
        database=spec.database,
        schema="MODELS",
        auto_create_table=True,
        overwrite=True,
        quote_identifiers=False,
    )
    return fqn


def write_diagnostics(
    session,
    spec: UseCaseSpec,
    results: list[EstimateResult],
    table: str = "DIAGNOSTICS",
) -> str:
    """Persist one row per estimator so the receipt and app can read them."""
    rows = []
    for r in results:
        rows.append(
            {
                "METHOD": r.method,
                "OUTCOME_COLUMN": r.outcome_column,
                "ESTIMATE": float(r.estimate),
                "CI_LOW": None if r.ci_low is None else float(r.ci_low),
                "CI_HIGH": None if r.ci_high is None else float(r.ci_high),
                "TARGET": str(r.diagnostics.get("target", "ate")),
                "IDENTIFIED": bool(r.identified),
                "N": int(r.n),
                "N_TREATED": int(r.n_treated),
                "CATE_SPREAD": None if r.cate_spread is None else float(r.cate_spread),
                "VERDICT": r.direction_verdict(spec.outcome),
                "DIAGNOSTICS_JSON": json.dumps(r.diagnostics, default=float),
                "NOTES": " | ".join(r.notes),
            }
        )
    out = pd.DataFrame(rows)
    session.write_pandas(
        out,
        table,
        database=spec.database,
        schema="MODELS",
        auto_create_table=True,
        overwrite=True,
        quote_identifiers=False,
    )
    return spec.models(table)


# ---------------------------------------------------------------------------
# internals
# ---------------------------------------------------------------------------


def _cross_fit_outcome_models(
    X: pd.DataFrame,
    t: np.ndarray,
    y: np.ndarray,
    outcome: Outcome,
    n_folds: int,
    random_state: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Out-of-fold predictions of the treated and untreated response surfaces."""
    mu1 = np.zeros(len(y), dtype=float)
    mu0 = np.zeros(len(y), dtype=float)
    splitter = StratifiedKFold(
        n_splits=n_folds, shuffle=True, random_state=random_state
    )

    for train_idx, test_idx in splitter.split(X, t):
        for arm, target in ((1, mu1), (0, mu0)):
            arm_rows = train_idx[t[train_idx] == arm]
            model = _outcome_model(outcome, random_state)
            model.fit(X.iloc[arm_rows], y[arm_rows])
            target[test_idx] = _predict(model, X.iloc[test_idx])
    return mu1, mu0


def _outcome_model(outcome: Outcome, random_state: int):
    if outcome.is_binary:
        return HistGradientBoostingClassifier(
            max_iter=200, learning_rate=0.08, max_depth=6, random_state=random_state
        )
    return HistGradientBoostingRegressor(
        max_iter=200, learning_rate=0.08, max_depth=6, random_state=random_state
    )


def _predict(model, X: pd.DataFrame) -> np.ndarray:
    if hasattr(model, "predict_proba"):
        return model.predict_proba(X)[:, 1]
    return model.predict(X)


def _pre_trend_test(event_study: pd.DataFrame) -> dict[str, float]:
    """Regress the treated-minus-control gap on time, using pre-periods only."""
    import statsmodels.formula.api as smf

    pre = event_study[event_study["IS_PRE"].astype(bool)].copy()
    fit = smf.ols("GAP ~ PERIOD_INDEX", data=pre).fit()
    return {
        "slope": float(fit.params["PERIOD_INDEX"]),
        "p_value": float(fit.pvalues["PERIOD_INDEX"]),
    }