Formative

from formative.causal import DAG, OLSObservational
from formative.game import minimax, maximax, maximin


# ---- Causal estimation (with heterogeneous effects) ----

dag = DAG()
dag.assume("ability").causes("education", "income")
dag.assume("segment").causes("income")
dag.assume("education").causes("income")

result = OLSObservational(
    dag,
    treatment="education",
    outcome="income",
    effect_modifier="segment",     # effect estimated per segment
).fit(df)
print(result.summary())


# ---- Refutation ----

print(result.refute(df).summary())


# ---- Decisions ----

decision = result.decide(cost=2.5, benefit=1.0)
print(decision)

for level, d in result.decide_by_group(cost=2.5, benefit=1.0).items():
    print(level, "→", d.optimal)   # different segments


# ---- Decisions via game theory ----

outcomes = decision.to_outcomes()
print(minimax(outcomes).solve())   # pessimistic choice
print(maximax(outcomes).solve())   # high risk, high reward
print(maximin(outcomes).solve())   # best worst-case