{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-001",
   "metadata": {},
   "source": [
    "# Module 5 — Lab 5: Interpreting Trees and Controlling Overfitting\n",
    "\n",
    "  \n",
    "**Topics:** tree rules, complexity, feature importance, train-test gap\n",
    "\n",
    "> **Solution version:** Run the cells in order, compare the outputs with your work, and read the suggested answers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-002",
   "metadata": {},
   "source": [
    "## Learning objectives\n",
    "\n",
    "After completing this lab, you will be able to:\n",
    "\n",
    "- compare bounded and unrestricted trees;\n",
    "- quantify depth, number of leaves, and train-test performance gaps;\n",
    "- visualize an explainable tree;\n",
    "- export root-to-leaf logic as text rules;\n",
    "- interpret feature importance cautiously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-003",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "DATA_PATH = Path(\"data/malware_features.csv\")\n",
    "\n",
    "if not DATA_PATH.exists():\n",
    "    raise FileNotFoundError(\n",
    "        \"Missing data/malware_features.csv. Follow data/README.md, place the three \"\n",
    "        \"official UCI CSV files under data/raw/, and run: python data/prepare_dataset.py\"\n",
    "    )\n",
    "\n",
    "\n",
    "from sklearn.dummy import DummyClassifier\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score, accuracy_score\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.tree import DecisionTreeClassifier\n",
    "\n",
    "\n",
    "from sklearn.tree import export_text, plot_tree\n",
    "\n",
    "\n",
    "def security_metrics(y_true, y_pred):\n",
    "    tn, fp, fn, tp = confusion_matrix(y_true, y_pred, labels=[0, 1]).ravel()\n",
    "    return {\n",
    "        \"accuracy\": accuracy_score(y_true, y_pred),\n",
    "        \"malware_precision\": precision_score(y_true, y_pred, zero_division=0),\n",
    "        \"malware_recall\": recall_score(y_true, y_pred, zero_division=0),\n",
    "        \"malware_f1\": f1_score(y_true, y_pred, zero_division=0),\n",
    "        \"false_positives\": int(fp),\n",
    "        \"false_negatives\": int(fn),\n",
    "        \"false_positive_rate\": fp / (fp + tn) if fp + tn else 0.0,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-004",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(DATA_PATH, low_memory=False)\n",
    "feature_cols = [column for column in df.columns if column not in {\"source\", \"label\"}]\n",
    "X = df[feature_cols].apply(pd.to_numeric, errors=\"coerce\")\n",
    "y = df[\"label\"].astype(int)\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X, y, test_size=0.20, stratify=y, random_state=42\n",
    ")\n",
    "\n",
    "def make_tree(max_depth, min_samples_leaf):\n",
    "    return Pipeline([\n",
    "        (\"imputer\", SimpleImputer(strategy=\"median\", add_indicator=True)),\n",
    "        (\"tree\", DecisionTreeClassifier(\n",
    "            max_depth=max_depth, min_samples_leaf=min_samples_leaf, random_state=42\n",
    "        )),\n",
    "    ])\n",
    "\n",
    "bounded_tree = make_tree(max_depth=4, min_samples_leaf=10)\n",
    "unrestricted_tree = make_tree(max_depth=None, min_samples_leaf=1)\n",
    "bounded_tree.fit(X_train, y_train)\n",
    "unrestricted_tree.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-005",
   "metadata": {},
   "source": [
    "## 1. Compare model complexity and generalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-006",
   "metadata": {},
   "outputs": [],
   "source": [
    "comparison_rows = []\n",
    "for name, model in {\"bounded\": bounded_tree, \"unrestricted\": unrestricted_tree}.items():\n",
    "    estimator = model.named_steps[\"tree\"]\n",
    "    train_metrics = security_metrics(y_train, model.predict(X_train))\n",
    "    test_metrics = security_metrics(y_test, model.predict(X_test))\n",
    "    comparison_rows.append({\n",
    "        \"model\": name,\n",
    "        \"depth\": estimator.get_depth(),\n",
    "        \"leaves\": estimator.get_n_leaves(),\n",
    "        \"train_f1\": train_metrics[\"malware_f1\"],\n",
    "        \"test_f1\": test_metrics[\"malware_f1\"],\n",
    "        \"f1_gap\": train_metrics[\"malware_f1\"] - test_metrics[\"malware_f1\"],\n",
    "        \"test_recall\": test_metrics[\"malware_recall\"],\n",
    "        \"test_fpr\": test_metrics[\"false_positive_rate\"],\n",
    "    })\n",
    "complexity_comparison = pd.DataFrame(comparison_rows).set_index(\"model\")\n",
    "display(complexity_comparison.round(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-007",
   "metadata": {},
   "source": [
    "## 2. Visualize and export the bounded tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-008",
   "metadata": {},
   "outputs": [],
   "source": [
    "imputer = bounded_tree.named_steps[\"imputer\"]\n",
    "estimator = bounded_tree.named_steps[\"tree\"]\n",
    "transformed_names = imputer.get_feature_names_out(feature_cols)\n",
    "\n",
    "plt.figure(figsize=(24, 12))\n",
    "plot_tree(\n",
    "    estimator,\n",
    "    feature_names=transformed_names,\n",
    "    class_names=[\"benign\", \"malware\"],\n",
    "    filled=True,\n",
    "    rounded=True,\n",
    "    proportion=True,\n",
    "    fontsize=7,\n",
    ")\n",
    "plt.title(\"Bounded malware Decision Tree\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "rule_text = export_text(estimator, feature_names=list(transformed_names), max_depth=4)\n",
    "print(rule_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-009",
   "metadata": {},
   "source": [
    "## 3. Inspect impurity-based feature importance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-010",
   "metadata": {},
   "outputs": [],
   "source": [
    "importance = pd.Series(\n",
    "    estimator.feature_importances_, index=transformed_names, name=\"importance\"\n",
    ").sort_values(ascending=False)\n",
    "top_importance = importance.head(12)\n",
    "display(top_importance.to_frame())\n",
    "\n",
    "ax = top_importance.sort_values().plot(kind=\"barh\", figsize=(9, 6), color=\"#4C78A8\")\n",
    "ax.set_title(\"Top bounded-tree feature importances\")\n",
    "ax.set_xlabel(\"Mean decrease in impurity\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-011",
   "metadata": {},
   "source": [
    "## Interpretation warning\n",
    "\n",
    "A rule describes this fitted tree, not a universal malware signature. Correlated features can substitute for one another, high-cardinality features may receive inflated impurity importance, and attackers may deliberately modify static properties. Treat rules as analyst evidence that requires validation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-012",
   "metadata": {},
   "source": [
    "## Self-check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-013",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert bounded_tree.named_steps[\"tree\"].get_depth() <= 4\n",
    "assert unrestricted_tree.named_steps[\"tree\"].get_depth() >= bounded_tree.named_steps[\"tree\"].get_depth()\n",
    "assert complexity_comparison.loc[\"unrestricted\", \"leaves\"] >= complexity_comparison.loc[\"bounded\", \"leaves\"]\n",
    "assert isinstance(rule_text, str) and \"class:\" in rule_text\n",
    "assert np.isclose(importance.sum(), 1.0)\n",
    "assert len(top_importance) <= 12\n",
    "print(\"✓ Lab 5 completed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-014",
   "metadata": {},
   "source": [
    "## Suggested answers\n",
    "\n",
    "1. **What indicates overfitting?** The unrestricted tree has more depth/leaves and typically a larger train-test performance gap.\n",
    "2. **How should a rule be written?** Combine each split along one root-to-leaf path with `AND`, then attach the leaf prediction and sample support.\n",
    "3. **Why is importance not causality?** It measures impurity reduction inside one fitted model and can reflect correlation or collection artifacts."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}