{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-000",
   "metadata": {},
   "source": [
    "# Lesson 6 — Rule Extraction & Interpretability\n",
    "\n",
    "**Companion notebook** for the Lesson 6 slides. You'll fit surrogate decision trees\n",
    "that imitate Lesson 5's clusters, choose how complex they may be, export the rules, and\n",
    "turn them into risk tiers that you test against next year's cost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-001",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check(label, got, expected, tol=1e-6):\n",
    "    \"\"\"Compare your answer with the expected value; prints ✅ or ❌ (never raises).\"\"\"\n",
    "    import math\n",
    "    try:\n",
    "        if isinstance(expected, float) or isinstance(got, float):\n",
    "            ok = math.isclose(float(got), float(expected), rel_tol=tol, abs_tol=tol)\n",
    "        elif isinstance(expected, (list, tuple)) and expected and isinstance(expected[0], str):\n",
    "            ok = list(got) == list(expected)\n",
    "        elif hasattr(expected, \"__len__\") and not isinstance(expected, str):\n",
    "            import numpy as np\n",
    "            ok = np.allclose(np.asarray(got, dtype=float), np.asarray(expected, dtype=float), rtol=tol, atol=tol)\n",
    "        else:\n",
    "            ok = got == expected\n",
    "    except Exception:\n",
    "        ok = False\n",
    "    print((\"✅ \" if ok else \"❌ \") + f\"{label}: got {got!r}, expected {expected!r}\")\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.cluster import KMeans\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-002",
   "metadata": {},
   "source": [
    "## 1. Rebuild Lesson 5's profiles and clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-003",
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = pd.read_parquet(\"../data/health/encounters.parquet\")\n",
    "mem = pd.read_csv(\"../data/health/members.csv\")\n",
    "e = enc[enc[\"service_date\"].dt.year == 2024]\n",
    "counts = e.pivot_table(index=\"member_id\", columns=\"service_category\", values=\"encounter_id\",\n",
    "                       aggfunc=\"count\", fill_value=0)\n",
    "counts.columns = [c.lower().replace(\" \", \"_\") for c in counts.columns]\n",
    "cols = [\"office_visit\", \"specialist\", \"lab\", \"imaging\", \"pharmacy\", \"emergency\",\n",
    "        \"outpatient_surgery\", \"inpatient\"]\n",
    "prof = counts.reindex(mem[\"member_id\"], fill_value=0)[cols]\n",
    "prof[\"total_allowed\"] = e.groupby(\"member_id\")[\"allowed\"].sum().reindex(prof.index, fill_value=0)\n",
    "prof[\"active_months\"] = (e.assign(m=e[\"service_date\"].dt.month).groupby(\"member_id\")[\"m\"]\n",
    "                          .nunique().reindex(prof.index, fill_value=0))\n",
    "\n",
    "X = StandardScaler().fit_transform(np.log1p(prof))\n",
    "lab = pd.Series(KMeans(3, n_init=10, random_state=0).fit(X).labels_, index=prof.index)\n",
    "order = prof.groupby(lab)[\"total_allowed\"].mean().sort_values().index\n",
    "cluster = lab.map(dict(zip(order, [\"Healthy\", \"Chronic-Maintenance\", \"Heavy-Use\"])))\n",
    "cluster.value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-004",
   "metadata": {},
   "source": [
    "## 2. A surrogate tree\n",
    "\n",
    "The surrogate learns the **cluster label** from the profile in raw units (counts,\n",
    "dollars, months), so its rules read naturally. Fidelity = agreement with the clusters\n",
    "on members the tree hasn't seen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-005",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_tr, X_te, y_tr, y_te = train_test_split(prof, cluster, test_size=0.3, random_state=0,\n",
    "                                          stratify=cluster)\n",
    "tree = DecisionTreeClassifier(max_depth=2, random_state=0).fit(X_tr, y_tr)\n",
    "print(f\"depth-2 fidelity: training {tree.score(X_tr, y_tr):.4f}, test {tree.score(X_te, y_te):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-006",
   "metadata": {},
   "source": [
    "---\n",
    "## Your turn 1 — how deep?\n",
    "\n",
    "For `max_depth` in 1, 2, 3, 4, 6, 10, and `None`, record the number of leaves\n",
    "(`get_n_leaves()`), training fidelity, and test fidelity in a DataFrame `depths`\n",
    "indexed by depth (use the string `\"None\"` for the unlimited tree)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-007",
   "metadata": {},
   "outputs": [],
   "source": [
    "depths = None\n",
    "...  # your code here\n",
    "check(\"depth-2 test fidelity\", None if depths is None else round(depths.loc[\"2\", \"test\"], 4), 0.9625)\n",
    "check(\"leaves, unlimited tree\", None if depths is None else int(depths.loc[\"None\", \"leaves\"]), 138)\n",
    "check(\"training fidelity, unlimited tree\", None if depths is None else round(depths.loc[\"None\", \"train\"], 4), 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-008",
   "metadata": {},
   "source": [
    "## 3. Cost-complexity pruning\n",
    "\n",
    "Grow a full tree, then prune branches worth less than α per leaf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-009",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = DecisionTreeClassifier(random_state=0).cost_complexity_pruning_path(X_tr, y_tr)\n",
    "res = []\n",
    "for a in path.ccp_alphas[::5]:\n",
    "    t = DecisionTreeClassifier(random_state=0, ccp_alpha=a).fit(X_tr, y_tr)\n",
    "    res.append((t.get_n_leaves(), t.score(X_tr, y_tr), t.score(X_te, y_te)))\n",
    "res = pd.DataFrame(res, columns=[\"leaves\", \"train\", \"test\"])\n",
    "plt.figure(figsize=(6.5, 3.2))\n",
    "plt.plot(res[\"leaves\"], res[\"train\"], \"o-\", label=\"training\")\n",
    "plt.plot(res[\"leaves\"], res[\"test\"], \"o-\", label=\"test\")\n",
    "plt.xscale(\"log\"); plt.xlabel(\"leaves\"); plt.ylabel(\"fidelity\"); plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-010",
   "metadata": {},
   "source": [
    "## Your turn 2 — export the rules\n",
    "\n",
    "Print the depth-2 tree with `export_text(tree, feature_names=list(prof.columns),\n",
    "decimals=1, show_weights=True)`, and draw it with `plot_tree`. Store the feature and\n",
    "threshold of the **first** split (`tree.tree_.feature[0]`, `tree.tree_.threshold[0]`)\n",
    "and the threshold on `total_allowed`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-011",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_feature = first_threshold = cost_threshold = None\n",
    "...  # your code here\n",
    "check(\"first split feature\", first_feature, \"active_months\")\n",
    "check(\"first split threshold\", None if first_threshold is None else round(first_threshold, 1), 8.5)\n",
    "check(\"cost threshold\", None if cost_threshold is None else round(cost_threshold, 1), 25437.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-012",
   "metadata": {},
   "source": [
    "## Your turn 3 — risk tiers, tested on next year\n",
    "\n",
    "Map every member to a tier with the depth-2 tree:\n",
    "Healthy → \"Tier 1 · Preferred\", Chronic-Maintenance → \"Tier 2 · Care management\",\n",
    "Heavy-Use → \"Tier 3 · Complex care\". Then tabulate members, mean 2024 cost, mean 2025\n",
    "cost, and the 2025 relativity (tier mean ÷ overall mean) by tier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-013",
   "metadata": {},
   "outputs": [],
   "source": [
    "TIERS = {\"Healthy\": \"Tier 1 · Preferred\", \"Chronic-Maintenance\": \"Tier 2 · Care management\",\n",
    "         \"Heavy-Use\": \"Tier 3 · Complex care\"}\n",
    "cost25 = (enc[enc[\"service_date\"].dt.year == 2025].groupby(\"member_id\")[\"allowed\"].sum()\n",
    "          .reindex(prof.index, fill_value=0))\n",
    "tiers = None\n",
    "...  # your code here\n",
    "check(\"Tier 3 members\", None if tiers is None else int(tiers.loc[\"Tier 3 · Complex care\", \"members\"]), 657)\n",
    "check(\"Tier 3 relativity (2025)\",\n",
    "      None if tiers is None else round(tiers.loc[\"Tier 3 · Complex care\", \"relativity_2025\"], 2), 4.09)\n",
    "check(\"Tier 1 mean 2025 cost\",\n",
    "      None if tiers is None else round(tiers.loc[\"Tier 1 · Preferred\", \"mean_2025\"]), 3874)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}