{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a3a90183",
   "metadata": {},
   "source": [
    "# Module 3 — Lab 5: A Controlled Class-Imbalance Experiment\n",
    "\n",
    " \n",
    "**Topics:** class imbalance, class weights, controlled comparison, metric trade-offs\n",
    "\n",
    "> **Solution version:** This notebook contains the complete solution. Run the cells in order and compare the results with your Lab notebook.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d80ad41",
   "metadata": {},
   "source": [
    "## Learning objectives\n",
    "\n",
    "After completing this lab, you will be able to:\n",
    "\n",
    "- measure the class imbalance in a security dataset;\n",
    "- compare two runs of the same model on the same split;\n",
    "- use `class_weight=\"balanced\"` without changing the rest of the workflow;\n",
    "- select a model configuration using attack precision, recall, and F1.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a4b369",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "DATA_PATH = Path(\"data/synthetic_security_flows.csv\")\n",
    "RANDOM_STATE = 42\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1301b149",
   "metadata": {},
   "source": [
    "## 1. Prepare one fixed experiment\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5ea362d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(DATA_PATH).drop_duplicates()\n",
    "df.columns = df.columns.str.strip()\n",
    "df = df.replace([np.inf, -np.inf], np.nan)\n",
    "\n",
    "X = df.drop(columns=\"Label\").select_dtypes(include=\"number\")\n",
    "y = df[\"Label\"].astype(str).str.strip().str.upper().ne(\"BENIGN\").astype(int)\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X, y, test_size=0.2, random_state=RANDOM_STATE, stratify=y\n",
    ")\n",
    "\n",
    "benign_count = int((y_train == 0).sum())\n",
    "attack_count = int((y_train == 1).sum())\n",
    "imbalance_ratio = benign_count / attack_count\n",
    "\n",
    "print(f\"BENIGN={benign_count}, ATTACK={attack_count}\")\n",
    "print(f\"BENIGN-to-ATTACK ratio: {imbalance_ratio:.2f}:1\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3dfcb43",
   "metadata": {},
   "source": [
    "## 2. Define a reusable training function\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51c588dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_measure(class_weight):\n",
    "    model = Pipeline([\n",
    "        (\"imputer\", SimpleImputer(strategy=\"median\")),\n",
    "        (\"model\", RandomForestClassifier(\n",
    "            n_estimators=100,\n",
    "            max_depth=12,\n",
    "            class_weight=class_weight,\n",
    "            random_state=RANDOM_STATE,\n",
    "            n_jobs=-1,\n",
    "        )),\n",
    "    ])\n",
    "    model.fit(X_train, y_train)\n",
    "    predictions = model.predict(X_test)\n",
    "    tn, fp, fn, tp = confusion_matrix(y_test, predictions).ravel()\n",
    "    return {\n",
    "        \"precision\": precision_score(y_test, predictions, zero_division=0),\n",
    "        \"recall\": recall_score(y_test, predictions, zero_division=0),\n",
    "        \"f1\": f1_score(y_test, predictions, zero_division=0),\n",
    "        \"false_positives\": int(fp),\n",
    "        \"false_negatives\": int(fn),\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c81e9645",
   "metadata": {},
   "source": [
    "## 3. Run the controlled comparison\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac3530d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_metrics = train_and_measure(class_weight=None)\n",
    "balanced_metrics = train_and_measure(class_weight=\"balanced\")\n",
    "\n",
    "comparison_df = pd.DataFrame(\n",
    "    [baseline_metrics, balanced_metrics],\n",
    "    index=[\"Baseline\", \"Balanced\"],\n",
    ")\n",
    "\n",
    "comparison_df.round(3)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4f99ffd",
   "metadata": {},
   "source": [
    "## 4. Make an evidence-based recommendation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc7c44fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "recommended_configuration = comparison_df[\"f1\"].idxmax()\n",
    "\n",
    "print(\"Recommended configuration:\", recommended_configuration)\n",
    "print(comparison_df.loc[recommended_configuration].round(3))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6e2a54c",
   "metadata": {},
   "source": [
    "## Self-check\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6389320",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert benign_count > attack_count\n",
    "assert imbalance_ratio > 1\n",
    "assert list(comparison_df.index) == [\"Baseline\", \"Balanced\"]\n",
    "assert {\"precision\", \"recall\", \"f1\", \"false_positives\", \"false_negatives\"}.issubset(comparison_df.columns)\n",
    "assert comparison_df[[\"precision\", \"recall\", \"f1\"]].apply(lambda col: col.between(0, 1).all()).all()\n",
    "assert recommended_configuration in {\"Baseline\", \"Balanced\"}\n",
    "print(\"✓ Lab 5 completed\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69c051ee",
   "metadata": {},
   "source": [
    "## Analyst questions\n",
    "\n",
    "1. Did class weighting improve attack recall in this run?\n",
    "2. What happened to false positives and false negatives?\n",
    "3. Would you choose the same configuration if analyst time were extremely limited?\n",
    "4. Why is changing only one variable important in a controlled experiment?\n"
   ]
  }
 ],
 "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
}