{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m15-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 15.1 · From Notebook to Decision Engine: Serving a Causal Model with FastAPI\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Train a CATE model for a retraining subsidy and save it as a versioned artifact\n",
    "- Wrap the model in a FastAPI endpoint that returns an effect estimate and a recommendation\n",
    "- Validate inputs with Pydantic and make the decision rule explicit\n",
    "- Test the API automatically with FastAPI's TestClient and pytest\n",
    "\n",
    "> Read the full explanations on the course website, then run each cell here in order (**Runtime → Run all** works on a fresh runtime)."
   ]
  },
  {
   "id": "m15-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why deploy a causal model?"
   ]
  },
  {
   "id": "m15-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas scikit-learn joblib fastapi httpx pytest\n",
    "import json, os, numpy as np, pandas as pd, joblib\n",
    "from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier\n",
    "rng = np.random.default_rng(15)\n",
    "os.makedirs(\"backend/models\", exist_ok=True)"
   ]
  },
  {
   "id": "m15-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and save the model"
   ]
  },
  {
   "id": "m15-l1-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the expected wage gain from a retraining subsidy for a given applicant?"
   ]
  },
  {
   "id": "m15-l1-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "n = 12_000\n",
    "X = pd.DataFrame({\n",
    "    \"poverty_index\": rng.normal(0, 1, n),\n",
    "    \"years_education\": rng.integers(8, 21, n),\n",
    "    \"unemployment_duration_months\": rng.gamma(2, 4, n).round(),\n",
    "})\n",
    "tau = 300 + 120 * np.minimum(X.unemployment_duration_months, 24) - 80 * (X.years_education - 12) + 300 * X.poverty_index\n",
    "p = 1 / (1 + np.exp(-(-0.3 + 0.05 * X.unemployment_duration_months + 0.3 * X.poverty_index)))\n",
    "d = rng.binomial(1, p)\n",
    "y = 28_000 - 2_000 * X.poverty_index + 900 * (X.years_education - 12) + d * tau + rng.normal(0, 2_000, n)\n",
    "\n",
    "# Doubly robust scores with 2-fold cross-fitting\n",
    "folds = rng.integers(0, 2, n); psi = np.zeros(n)\n",
    "for k in (0, 1):\n",
    "    tr, te = folds != k, folds == k\n",
    "    mu1 = GradientBoostingRegressor().fit(X[tr & (d == 1)], y[tr & (d == 1)])\n",
    "    mu0 = GradientBoostingRegressor().fit(X[tr & (d == 0)], y[tr & (d == 0)])\n",
    "    e = np.clip(GradientBoostingClassifier().fit(X[tr], d[tr]).predict_proba(X[te])[:, 1], 0.02, 0.98)\n",
    "    m1, m0 = mu1.predict(X[te]), mu0.predict(X[te])\n",
    "    psi[te] = m1 - m0 + d[te] * (y[te] - m1) / e - (1 - d[te]) * (y[te] - m0) / (1 - e)\n",
    "\n",
    "cate_model = GradientBoostingRegressor(n_estimators=200, max_depth=2, learning_rate=0.05).fit(X, psi)\n",
    "print(f\"Estimated ATE {psi.mean():,.0f} (true {tau.mean():,.0f}); \"\n",
    "      f\"corr(predicted CATE, true CATE) = {np.corrcoef(cate_model.predict(X), tau)[0, 1]:.2f}\")\n",
    "\n",
    "artifact = {\"model\": cate_model, \"features\": list(X.columns), \"version\": \"2026.11-dr-gbm\",\n",
    "            \"ate\": float(psi.mean()), \"ate_se\": float(psi.std() / np.sqrt(n))}\n",
    "joblib.dump(artifact, \"backend/models/cate_model.joblib\")\n",
    "print(\"Saved backend/models/cate_model.joblib\")"
   ]
  },
  {
   "id": "m15-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The FastAPI decision engine"
   ]
  },
  {
   "id": "m15-l1-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "%%writefile backend/app.py\n",
    "from pathlib import Path\n",
    "import joblib\n",
    "import pandas as pd\n",
    "from fastapi import FastAPI\n",
    "from fastapi.middleware.cors import CORSMiddleware\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "ARTIFACT = joblib.load(Path(__file__).parent / \"models\" / \"cate_model.joblib\")\n",
    "SUBSIDY_COST_CAD = 1_200.0\n",
    "\n",
    "app = FastAPI(title=\"Causal Decision Engine\", version=ARTIFACT[\"version\"])\n",
    "app.add_middleware(CORSMiddleware, allow_origins=[\"*\"], allow_methods=[\"*\"], allow_headers=[\"*\"])\n",
    "\n",
    "\n",
    "class CitizenProfile(BaseModel):\n",
    "    poverty_index: float = Field(..., ge=-4, le=4, description=\"Standardised poverty index\")\n",
    "    years_education: int = Field(..., ge=0, le=25)\n",
    "    unemployment_duration_months: int = Field(..., ge=0, le=120)\n",
    "\n",
    "\n",
    "@app.get(\"/health\")\n",
    "def health():\n",
    "    return {\"status\": \"ok\", \"model_version\": ARTIFACT[\"version\"]}\n",
    "\n",
    "\n",
    "@app.post(\"/predict_impact\")\n",
    "def predict_impact(profile: CitizenProfile):\n",
    "    row = pd.DataFrame([profile.model_dump()])[ARTIFACT[\"features\"]]\n",
    "    lift = float(ARTIFACT[\"model\"].predict(row)[0])\n",
    "    return {\n",
    "        \"expected_wage_lift_cad\": round(lift, 2),\n",
    "        \"subsidy_cost_cad\": SUBSIDY_COST_CAD,\n",
    "        \"net_benefit_cad\": round(lift - SUBSIDY_COST_CAD, 2),\n",
    "        \"recommend_intervention\": bool(lift > SUBSIDY_COST_CAD),\n",
    "        \"average_effect_cad\": round(ARTIFACT[\"ate\"], 2),\n",
    "        \"model_version\": ARTIFACT[\"version\"],\n",
    "    }"
   ]
  },
  {
   "id": "m15-l1-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automated tests"
   ]
  },
  {
   "id": "m15-l1-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How do I make sure the API keeps returning sensible, well-formed decisions after every change?"
   ]
  },
  {
   "id": "m15-l1-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "%%writefile tests/test_api.py\n",
    "from fastapi.testclient import TestClient\n",
    "from backend.app import app\n",
    "\n",
    "client = TestClient(app)\n",
    "\n",
    "\n",
    "def test_health():\n",
    "    assert client.get(\"/health\").json()[\"status\"] == \"ok\"\n",
    "\n",
    "\n",
    "def test_predict_impact_returns_valid_recommendation():\n",
    "    payload = {\"poverty_index\": 1.5, \"years_education\": 10, \"unemployment_duration_months\": 18}\n",
    "    r = client.post(\"/predict_impact\", json=payload)\n",
    "    assert r.status_code == 200\n",
    "    data = r.json()\n",
    "    assert isinstance(data[\"recommend_intervention\"], bool)\n",
    "    assert -10_000 < data[\"expected_wage_lift_cad\"] < 50_000      # no impossible effects\n",
    "\n",
    "\n",
    "def test_long_term_unemployed_gain_more():\n",
    "    base = {\"poverty_index\": 0.0, \"years_education\": 12}\n",
    "    short = client.post(\"/predict_impact\", json=base | {\"unemployment_duration_months\": 1}).json()\n",
    "    long = client.post(\"/predict_impact\", json=base | {\"unemployment_duration_months\": 20}).json()\n",
    "    assert long[\"expected_wage_lift_cad\"] > short[\"expected_wage_lift_cad\"]\n",
    "\n",
    "\n",
    "def test_rejects_invalid_input():\n",
    "    r = client.post(\"/predict_impact\", json={\"poverty_index\": 0, \"years_education\": -3, \"unemployment_duration_months\": 5})\n",
    "    assert r.status_code == 422"
   ]
  },
  {
   "id": "m15-l1-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import pytest\n",
    "exit_code = pytest.main([\"-q\", \"tests/test_api.py\"])\n",
    "print(\"All tests passed\" if exit_code == 0 else f\"pytest exit code {exit_code}\")\n",
    "assert exit_code == 0"
   ]
  },
  {
   "id": "m15-l1-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Add a `/predict_batch` endpoint that accepts a list of profiles and a `budget_share` (for example 0.3) and\n",
    "returns the profiles ranked by expected net benefit, with `recommend_intervention` set only for the top\n",
    "share. Write one test for it."
   ]
  },
  {
   "id": "m15-l1-13",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m15-l1-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}