{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0",
   "mimetype": "text/x-python",
   "file_extension": ".py",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python"
  },
  "colab": {
   "provenance": [],
   "gpuType": "T4",
   "name": "RL Case Study: Adaptive Chemotherapy Dose Optimization -- Vizuara"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \ud83d\udd27 Setup: Run this cell first!\n",
    "# Check GPU availability and install dependencies\n",
    "\n",
    "import torch\n",
    "import sys\n",
    "\n",
    "# Check GPU\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "    print(f\"\u2705 GPU available: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "    print(\"\u26a0\ufe0f No GPU detected. Some cells may run slowly.\")\n",
    "    print(\"   Go to Runtime \u2192 Change runtime type \u2192 GPU\")\n",
    "\n",
    "print(f\"\\n\ud83d\udce6 Python {sys.version.split()[0]}\")\n",
    "print(f\"\ud83d\udd25 PyTorch {torch.__version__}\")\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "SEED = 42\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(SEED)\n",
    "\n",
    "print(f\"\ud83c\udfb2 Random seed set to {SEED}\")\n",
    "\n",
    "%matplotlib inline"
   ],
   "id": "setup_cell"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RL Case Study: Adaptive Chemotherapy Dose Optimization\n",
    "\n",
    "## NovaCure Therapeutics -- Personalizing Cancer Treatment with Q-Learning\n",
    "\n",
    "In this notebook, we implement the full case study: an RL agent that learns to optimize chemotherapy dosing for individual patients. We build a pharmacokinetic/pharmacodynamic (PK/PD) patient simulator, wrap it in a Gymnasium-compatible environment, and train a Q-learning agent to outperform fixed-dose and rule-based protocols."
   ],
   "id": "cell_1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import defaultdict\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "print(\"RL Case Study: Adaptive Chemotherapy Dose Optimization\")\n",
    "print(\"NovaCure Therapeutics -- Vizuara\")\n",
    "print(\"=\" * 55)"
   ],
   "id": "cell_2"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part A: Patient Simulator (PK/PD Model)\n",
    "\n",
    "We model tumor dynamics using the Gompertzian growth model and log-kill drug effect. Inter-patient variability is captured through randomized PK/PD parameters."
   ],
   "id": "cell_3"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PatientSimulator:\n",
    "    \"\"\"\n",
    "    Simulated patient for chemotherapy dose optimization.\n",
    "\n",
    "    Tumor dynamics: Gompertzian growth with log-kill drug effect\n",
    "    Toxicity: WBC suppression proportional to dose, with recovery\n",
    "\n",
    "    State: [tumor_size, wbc, toxicity_grade, kidney_function, cycle_number]\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, seed=None):\n",
    "        if seed is not None:\n",
    "            np.random.seed(seed)\n",
    "\n",
    "        # Patient-specific parameters (sampled for inter-patient variability)\n",
    "        self.tumor_growth_rate = np.random.uniform(0.02, 0.08)\n",
    "        self.tumor_carrying_capacity = np.random.uniform(80, 120)  # mm\n",
    "        self.drug_sensitivity = np.random.uniform(0.15, 0.35)\n",
    "        self.toxicity_sensitivity = np.random.uniform(0.3, 0.7)\n",
    "        self.wbc_recovery_rate = np.random.uniform(0.4, 0.8)\n",
    "        self.baseline_wbc = np.random.uniform(6000, 10000)  # cells/uL\n",
    "        self.baseline_egfr = np.random.uniform(80, 120)  # mL/min\n",
    "\n",
    "        # Initial state\n",
    "        self.tumor_size = np.random.uniform(30, 60)  # mm\n",
    "        self.wbc = self.baseline_wbc\n",
    "        self.toxicity_grade = 0\n",
    "        self.egfr = self.baseline_egfr\n",
    "        self.cycle = 1\n",
    "        self.max_cycles = 8\n",
    "\n",
    "        # Track cumulative drug exposure\n",
    "        self.cumulative_dose = 0.0\n",
    "\n",
    "    def get_state(self):\n",
    "        \"\"\"Return normalized state vector.\"\"\"\n",
    "        return np.array([\n",
    "            self.tumor_size / 100.0,          # normalize to ~[0, 1]\n",
    "            self.wbc / 10000.0,               # normalize to ~[0, 1]\n",
    "            self.toxicity_grade / 4.0,         # normalize to [0, 1]\n",
    "            self.egfr / 120.0,                 # normalize to ~[0, 1]\n",
    "            self.cycle / self.max_cycles,      # normalize to [0, 1]\n",
    "        ])\n",
    "\n",
    "    def step(self, dose_fraction):\n",
    "        \"\"\"\n",
    "        Simulate one treatment cycle.\n",
    "\n",
    "        Args:\n",
    "            dose_fraction: fraction of standard dose (e.g., 0.8 = 80% of standard)\n",
    "\n",
    "        Returns:\n",
    "            state, reward, done\n",
    "        \"\"\"\n",
    "        # --- Tumor dynamics (Gompertzian growth + log-kill) ---\n",
    "        # Natural growth\n",
    "        growth = self.tumor_growth_rate * self.tumor_size * np.log(\n",
    "            self.tumor_carrying_capacity / max(self.tumor_size, 1.0)\n",
    "        )\n",
    "\n",
    "        # Drug effect (log-kill: kills a fraction proportional to dose)\n",
    "        drug_kill = self.drug_sensitivity * dose_fraction * self.tumor_size\n",
    "\n",
    "        # Add noise\n",
    "        noise = np.random.normal(0, 0.5)\n",
    "\n",
    "        tumor_change = growth - drug_kill + noise\n",
    "        old_tumor = self.tumor_size\n",
    "        self.tumor_size = max(1.0, self.tumor_size + tumor_change)\n",
    "\n",
    "        # --- Hematological toxicity (WBC suppression) ---\n",
    "        # Drug suppresses WBC proportional to dose\n",
    "        wbc_suppression = self.toxicity_sensitivity * dose_fraction * self.wbc * 0.3\n",
    "        # Natural recovery toward baseline\n",
    "        wbc_recovery = self.wbc_recovery_rate * (self.baseline_wbc - self.wbc) * 0.2\n",
    "\n",
    "        self.wbc = max(500, self.wbc - wbc_suppression + wbc_recovery + np.random.normal(0, 200))\n",
    "\n",
    "        # --- Toxicity grading (based on WBC count, CTCAE-like) ---\n",
    "        if self.wbc >= 4000:\n",
    "            self.toxicity_grade = 0\n",
    "        elif self.wbc >= 3000:\n",
    "            self.toxicity_grade = 1\n",
    "        elif self.wbc >= 2000:\n",
    "            self.toxicity_grade = 2\n",
    "        elif self.wbc >= 1000:\n",
    "            self.toxicity_grade = 3\n",
    "        else:\n",
    "            self.toxicity_grade = 4\n",
    "\n",
    "        # --- Kidney function (gradual decline with cumulative dose) ---\n",
    "        self.cumulative_dose += dose_fraction\n",
    "        egfr_decline = 0.5 * dose_fraction + np.random.normal(0, 0.3)\n",
    "        self.egfr = max(30, self.egfr - egfr_decline)\n",
    "\n",
    "        # --- Reward computation ---\n",
    "        # Component 1: Tumor shrinkage (positive reward)\n",
    "        tumor_shrinkage = (old_tumor - self.tumor_size) / old_tumor  # fractional change\n",
    "        r_tumor = 10.0 * tumor_shrinkage\n",
    "\n",
    "        # Component 2: Toxicity penalty\n",
    "        r_toxicity = 0.0\n",
    "        if self.toxicity_grade >= 3:\n",
    "            r_toxicity = -5.0\n",
    "        if self.toxicity_grade >= 4:\n",
    "            r_toxicity = -15.0\n",
    "\n",
    "        # Component 3: Completion bonus\n",
    "        r_completion = 0.5 if self.toxicity_grade < 3 else 0.0\n",
    "\n",
    "        reward = r_tumor + r_toxicity + r_completion\n",
    "\n",
    "        # --- Episode termination ---\n",
    "        self.cycle += 1\n",
    "        done = False\n",
    "\n",
    "        # End conditions\n",
    "        if self.cycle > self.max_cycles:\n",
    "            done = True\n",
    "            # Bonus for completing all cycles\n",
    "            reward += 5.0\n",
    "        if self.toxicity_grade >= 4 and np.random.random() < 0.5:\n",
    "            done = True  # treatment discontinuation\n",
    "            reward -= 10.0\n",
    "        if self.tumor_size < 5.0:\n",
    "            done = True\n",
    "            reward += 20.0  # near-complete response\n",
    "\n",
    "        return self.get_state(), reward, done\n",
    "\n",
    "\n",
    "# Test the simulator\n",
    "patient = PatientSimulator(seed=42)\n",
    "print(\"Initial patient state:\")\n",
    "print(f\"  Tumor size: {patient.tumor_size:.1f} mm\")\n",
    "print(f\"  WBC count:  {patient.wbc:.0f} cells/uL\")\n",
    "print(f\"  Toxicity:   Grade {patient.toxicity_grade}\")\n",
    "print(f\"  Kidney (eGFR): {patient.egfr:.1f} mL/min\")\n",
    "print(f\"  Cycle: {patient.cycle}/{patient.max_cycles}\")\n",
    "\n",
    "# Run a few cycles with standard dose\n",
    "print(\"\\nSimulating 3 cycles at standard dose (1.0):\")\n",
    "for i in range(3):\n",
    "    state, reward, done = patient.step(1.0)\n",
    "    print(f\"  Cycle {i+2}: tumor={patient.tumor_size:.1f}mm, \"\n",
    "          f\"WBC={patient.wbc:.0f}, tox=Grade {patient.toxicity_grade}, \"\n",
    "          f\"reward={reward:.2f}, done={done}\")"
   ],
   "id": "cell_4"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part B: Gymnasium-Compatible Environment"
   ],
   "id": "cell_5"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChemoDoseEnv:\n",
    "    \"\"\"\n",
    "    Gymnasium-like environment for chemotherapy dose optimization.\n",
    "\n",
    "    Action space: 6 discrete dose levels\n",
    "        0: 0.6x standard dose\n",
    "        1: 0.7x standard dose\n",
    "        2: 0.8x standard dose\n",
    "        3: 0.9x standard dose\n",
    "        4: 1.0x standard dose (standard)\n",
    "        5: 1.1x standard dose\n",
    "\n",
    "    Observation: 5-dimensional normalized state\n",
    "    \"\"\"\n",
    "\n",
    "    DOSE_LEVELS = [0.6, 0.7, 0.8, 0.9, 1.0, 1.1]\n",
    "\n",
    "    def __init__(self):\n",
    "        self.n_actions = len(self.DOSE_LEVELS)\n",
    "        self.n_obs = 5\n",
    "        self.patient = None\n",
    "\n",
    "    def reset(self, seed=None):\n",
    "        self.patient = PatientSimulator(seed=seed)\n",
    "        return self.patient.get_state()\n",
    "\n",
    "    def step(self, action):\n",
    "        dose = self.DOSE_LEVELS[action]\n",
    "        state, reward, done = self.patient.step(dose)\n",
    "        info = {\n",
    "            'tumor_size': self.patient.tumor_size,\n",
    "            'wbc': self.patient.wbc,\n",
    "            'toxicity_grade': self.patient.toxicity_grade,\n",
    "            'egfr': self.patient.egfr,\n",
    "            'cycle': self.patient.cycle,\n",
    "            'dose': dose,\n",
    "        }\n",
    "        return state, reward, done, info\n",
    "\n",
    "\n",
    "# Test the environment\n",
    "env = ChemoDoseEnv()\n",
    "state = env.reset(seed=0)\n",
    "print(\"ChemoDoseEnv created.\")\n",
    "print(f\"  Action space: {env.n_actions} actions (dose levels: {env.DOSE_LEVELS})\")\n",
    "print(f\"  State dimension: {env.n_obs}\")\n",
    "print(f\"  Initial state: {state}\")"
   ],
   "id": "cell_6"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part C: Baseline Agents"
   ],
   "id": "cell_7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FixedDoseAgent:\n",
    "    \"\"\"Always prescribes the standard dose (1.0x).\"\"\"\n",
    "\n",
    "    def choose_action(self, state):\n",
    "        return 4  # index 4 = 1.0x standard dose\n",
    "\n",
    "    def name(self):\n",
    "        return \"Fixed Dose (1.0x)\"\n",
    "\n",
    "\n",
    "class RuleBasedAgent:\n",
    "    \"\"\"\n",
    "    Adjusts dose based on current toxicity grade.\n",
    "\n",
    "    - Grade 0-1: maintain or increase dose\n",
    "    - Grade 2: maintain dose\n",
    "    - Grade 3: reduce by 20%\n",
    "    - Grade 4: reduce to minimum\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.current_dose_idx = 4  # start at standard\n",
    "\n",
    "    def choose_action(self, state):\n",
    "        toxicity = state[2] * 4.0  # denormalize\n",
    "\n",
    "        if toxicity >= 4:\n",
    "            self.current_dose_idx = 0  # 0.6x\n",
    "        elif toxicity >= 3:\n",
    "            self.current_dose_idx = max(0, self.current_dose_idx - 2)  # reduce by 2 levels\n",
    "        elif toxicity >= 2:\n",
    "            self.current_dose_idx = max(0, self.current_dose_idx - 1)  # reduce by 1 level\n",
    "        elif toxicity <= 1:\n",
    "            self.current_dose_idx = min(5, self.current_dose_idx + 1)  # increase by 1 level\n",
    "\n",
    "        return self.current_dose_idx\n",
    "\n",
    "    def reset(self):\n",
    "        self.current_dose_idx = 4\n",
    "\n",
    "    def name(self):\n",
    "        return \"Rule-Based\"\n",
    "\n",
    "\n",
    "print(\"Baseline agents defined:\")\n",
    "print(\"  1. Fixed Dose: always prescribes 1.0x standard\")\n",
    "print(\"  2. Rule-Based: adjusts dose reactively based on toxicity grade\")"
   ],
   "id": "cell_8"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part D: Q-Learning Agent"
   ],
   "id": "cell_9"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChemoQLearningAgent:\n",
    "    \"\"\"\n",
    "    Q-learning agent for chemotherapy dose optimization.\n",
    "\n",
    "    Discretizes the 5-dimensional continuous state into bins\n",
    "    for table-based Q-learning.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, n_actions=6, n_bins=6, alpha=0.1, gamma=0.95,\n",
    "                 epsilon=1.0, epsilon_decay=0.999, epsilon_min=0.05):\n",
    "        self.n_actions = n_actions\n",
    "        self.n_bins = n_bins\n",
    "        self.alpha = alpha\n",
    "        self.gamma = gamma\n",
    "        self.epsilon = epsilon\n",
    "        self.epsilon_decay = epsilon_decay\n",
    "        self.epsilon_min = epsilon_min\n",
    "\n",
    "        self.q_table = defaultdict(lambda: np.zeros(n_actions))\n",
    "\n",
    "        # Bin edges for each state dimension\n",
    "        self.bins = [\n",
    "            np.linspace(0, 1, n_bins),      # tumor size (normalized)\n",
    "            np.linspace(0, 1.2, n_bins),     # WBC (normalized)\n",
    "            np.linspace(0, 1, n_bins),       # toxicity grade (normalized)\n",
    "            np.linspace(0.25, 1, n_bins),    # kidney function (normalized)\n",
    "            np.linspace(0, 1, n_bins),       # cycle number (normalized)\n",
    "        ]\n",
    "\n",
    "    def discretize(self, state):\n",
    "        discrete = []\n",
    "        for i, val in enumerate(state):\n",
    "            idx = np.digitize(val, self.bins[i])\n",
    "            discrete.append(idx)\n",
    "        return tuple(discrete)\n",
    "\n",
    "    def choose_action(self, state):\n",
    "        disc_state = self.discretize(state)\n",
    "\n",
    "        if np.random.random() < self.epsilon:\n",
    "            return np.random.randint(self.n_actions)\n",
    "        else:\n",
    "            return np.argmax(self.q_table[disc_state])\n",
    "\n",
    "    def update(self, state, action, reward, next_state, done):\n",
    "        s = self.discretize(state)\n",
    "        ns = self.discretize(next_state)\n",
    "\n",
    "        if done:\n",
    "            target = reward\n",
    "        else:\n",
    "            target = reward + self.gamma * np.max(self.q_table[ns])\n",
    "\n",
    "        self.q_table[s][action] += self.alpha * (target - self.q_table[s][action])\n",
    "\n",
    "    def decay_epsilon(self):\n",
    "        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)\n",
    "\n",
    "    def name(self):\n",
    "        return \"Q-Learning\"\n",
    "\n",
    "\n",
    "print(\"Q-Learning agent defined.\")\n",
    "print(\"Update rule: Q(s,a) <- Q(s,a) + alpha * [r + gamma * max Q(s',a') - Q(s,a)]\")"
   ],
   "id": "cell_10"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part E: Training"
   ],
   "id": "cell_11"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_agent(agent, n_episodes=200, verbose=False):\n",
    "    \"\"\"Evaluate an agent over many simulated patients.\"\"\"\n",
    "    env = ChemoDoseEnv()\n",
    "\n",
    "    results = {\n",
    "        'rewards': [],\n",
    "        'tumor_responses': [],\n",
    "        'toxicity_events': [],  # grade 3+\n",
    "        'completions': [],      # completed all 8 cycles\n",
    "        'final_tumor_sizes': [],\n",
    "        'dose_histories': [],\n",
    "    }\n",
    "\n",
    "    for ep in range(n_episodes):\n",
    "        state = env.reset(seed=ep + 10000)\n",
    "        total_reward = 0\n",
    "        tox_events = 0\n",
    "        doses = []\n",
    "\n",
    "        if hasattr(agent, 'reset'):\n",
    "            agent.reset()\n",
    "\n",
    "        for step in range(8):\n",
    "            if hasattr(agent, 'choose_action'):\n",
    "                action = agent.choose_action(state)\n",
    "            else:\n",
    "                action = agent(state)  # function-based agent\n",
    "\n",
    "            next_state, reward, done, info = env.step(action)\n",
    "            total_reward += reward\n",
    "            doses.append(info['dose'])\n",
    "\n",
    "            if info['toxicity_grade'] >= 3:\n",
    "                tox_events += 1\n",
    "\n",
    "            state = next_state\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        results['rewards'].append(total_reward)\n",
    "        results['final_tumor_sizes'].append(info['tumor_size'])\n",
    "        results['tumor_responses'].append(1 if info['tumor_size'] < 30 else 0)\n",
    "        results['toxicity_events'].append(tox_events)\n",
    "        results['completions'].append(1 if info['cycle'] > 8 else 0)\n",
    "        results['dose_histories'].append(doses)\n",
    "\n",
    "    return results\n",
    "\n",
    "\n",
    "# Train the Q-learning agent\n",
    "print(\"Training Q-learning agent...\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "env = ChemoDoseEnv()\n",
    "q_agent = ChemoQLearningAgent(\n",
    "    n_actions=6,\n",
    "    n_bins=6,\n",
    "    alpha=0.1,\n",
    "    gamma=0.95,\n",
    "    epsilon=1.0,\n",
    "    epsilon_decay=0.9995,\n",
    "    epsilon_min=0.05,\n",
    ")\n",
    "\n",
    "n_training_episodes = 20000\n",
    "training_rewards = []\n",
    "\n",
    "for ep in range(n_training_episodes):\n",
    "    state = env.reset(seed=ep)\n",
    "    total_reward = 0\n",
    "\n",
    "    for step in range(8):\n",
    "        action = q_agent.choose_action(state)\n",
    "        next_state, reward, done, info = env.step(action)\n",
    "        q_agent.update(state, action, reward, next_state, done)\n",
    "        state = next_state\n",
    "        total_reward += reward\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "    training_rewards.append(total_reward)\n",
    "    q_agent.decay_epsilon()\n",
    "\n",
    "    if (ep + 1) % 5000 == 0:\n",
    "        recent_avg = np.mean(training_rewards[-500:])\n",
    "        print(f\"  Episode {ep+1:>6} | Avg reward (last 500): {recent_avg:>7.2f} | \"\n",
    "              f\"Epsilon: {q_agent.epsilon:.4f} | Q-table size: {len(q_agent.q_table)}\")\n",
    "\n",
    "print(f\"\\nTraining complete. Q-table entries: {len(q_agent.q_table)}\")"
   ],
   "id": "cell_12"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot training curve\n",
    "fig, ax = plt.subplots(1, 1, figsize=(12, 5))\n",
    "window = 200\n",
    "smoothed = [np.mean(training_rewards[max(0,i-window):i+1]) for i in range(len(training_rewards))]\n",
    "ax.plot(smoothed, color='blue', alpha=0.8, linewidth=1.5)\n",
    "ax.set_xlabel('Training Episode', fontsize=12)\n",
    "ax.set_ylabel('Episode Reward (smoothed)', fontsize=12)\n",
    "ax.set_title('Q-Learning Agent: Training Progress', fontsize=14, fontweight='bold')\n",
    "ax.grid(True, alpha=0.3)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ],
   "id": "cell_13"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part F: Evaluation and Comparison"
   ],
   "id": "cell_14"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate all three agents\n",
    "print(\"Evaluating agents on 200 held-out patients...\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "fixed_agent = FixedDoseAgent()\n",
    "rule_agent = RuleBasedAgent()\n",
    "\n",
    "fixed_results = evaluate_agent(fixed_agent, n_episodes=200)\n",
    "rule_results = evaluate_agent(rule_agent, n_episodes=200)\n",
    "\n",
    "# For Q-agent, set epsilon to 0 (greedy)\n",
    "q_agent.epsilon = 0.0\n",
    "q_results = evaluate_agent(q_agent, n_episodes=200)\n",
    "\n",
    "# Print summary table\n",
    "agents = ['Fixed Dose', 'Rule-Based', 'Q-Learning']\n",
    "all_results = [fixed_results, rule_results, q_results]\n",
    "\n",
    "print(f\"\\n{'Metric':<30} | {'Fixed Dose':>12} | {'Rule-Based':>12} | {'Q-Learning':>12}\")\n",
    "print(\"-\" * 75)\n",
    "\n",
    "for metric, key, fmt in [\n",
    "    ('Avg Total Reward', 'rewards', '.2f'),\n",
    "    ('Tumor Response Rate', 'tumor_responses', '.1%'),\n",
    "    ('Avg Grade 3+ Tox Events', 'toxicity_events', '.2f'),\n",
    "    ('Treatment Completion Rate', 'completions', '.1%'),\n",
    "    ('Avg Final Tumor Size (mm)', 'final_tumor_sizes', '.1f'),\n",
    "]:\n",
    "    values = []\n",
    "    for res in all_results:\n",
    "        val = np.mean(res[key])\n",
    "        values.append(val)\n",
    "\n",
    "    if '%' in fmt:\n",
    "        print(f\"{metric:<30} | {values[0]:>11.1%} | {values[1]:>11.1%} | {values[2]:>11.1%}\")\n",
    "    else:\n",
    "        f = fmt.replace('.', ':>12.').replace('f', 'f')\n",
    "        print(f\"{metric:<30} | {values[0]:>12{fmt}} | {values[1]:>12{fmt}} | {values[2]:>12{fmt}}\")"
   ],
   "id": "cell_15"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Comprehensive visualization\n",
    "fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n",
    "\n",
    "# Plot 1: Reward distributions\n",
    "for i, (name, results, color) in enumerate(zip(\n",
    "    agents, all_results, ['red', 'orange', 'blue'])):\n",
    "    axes[0, 0].hist(results['rewards'], bins=20, alpha=0.5, label=name, color=color)\n",
    "axes[0, 0].set_xlabel('Total Episode Reward', fontsize=11)\n",
    "axes[0, 0].set_ylabel('Count', fontsize=11)\n",
    "axes[0, 0].set_title('Reward Distribution by Agent', fontsize=13, fontweight='bold')\n",
    "axes[0, 0].legend(fontsize=10)\n",
    "\n",
    "# Plot 2: Final tumor sizes\n",
    "for i, (name, results, color) in enumerate(zip(\n",
    "    agents, all_results, ['red', 'orange', 'blue'])):\n",
    "    axes[0, 1].hist(results['final_tumor_sizes'], bins=20, alpha=0.5, label=name, color=color)\n",
    "axes[0, 1].set_xlabel('Final Tumor Size (mm)', fontsize=11)\n",
    "axes[0, 1].set_ylabel('Count', fontsize=11)\n",
    "axes[0, 1].set_title('Final Tumor Size Distribution', fontsize=13, fontweight='bold')\n",
    "axes[0, 1].legend(fontsize=10)\n",
    "\n",
    "# Plot 3: Bar chart comparison\n",
    "metrics = ['Tumor\\nResponse', 'Treatment\\nCompletion', 'Low Toxicity\\n(Grade < 3)']\n",
    "metric_keys = ['tumor_responses', 'completions', 'toxicity_events']\n",
    "\n",
    "x = np.arange(len(metrics))\n",
    "width = 0.25\n",
    "\n",
    "for i, (name, results, color) in enumerate(zip(\n",
    "    agents, all_results, ['red', 'orange', 'blue'])):\n",
    "    vals = [\n",
    "        np.mean(results['tumor_responses']),\n",
    "        np.mean(results['completions']),\n",
    "        1.0 - np.mean([min(1, t) for t in results['toxicity_events']]),\n",
    "    ]\n",
    "    axes[1, 0].bar(x + i*width, vals, width, label=name, color=color, alpha=0.7)\n",
    "\n",
    "axes[1, 0].set_ylabel('Rate', fontsize=11)\n",
    "axes[1, 0].set_title('Clinical Outcomes Comparison', fontsize=13, fontweight='bold')\n",
    "axes[1, 0].set_xticks(x + width)\n",
    "axes[1, 0].set_xticklabels(metrics, fontsize=10)\n",
    "axes[1, 0].legend(fontsize=10)\n",
    "axes[1, 0].set_ylim(0, 1.1)\n",
    "\n",
    "# Plot 4: Average dosing trajectory\n",
    "for name, results, color in zip(agents, all_results, ['red', 'orange', 'blue']):\n",
    "    all_doses = results['dose_histories']\n",
    "    max_len = max(len(d) for d in all_doses)\n",
    "    avg_doses = []\n",
    "    for step in range(max_len):\n",
    "        step_doses = [d[step] for d in all_doses if len(d) > step]\n",
    "        avg_doses.append(np.mean(step_doses))\n",
    "    axes[1, 1].plot(range(1, len(avg_doses)+1), avg_doses, 'o-',\n",
    "                     color=color, label=name, linewidth=2, markersize=6)\n",
    "\n",
    "axes[1, 1].axhline(y=1.0, color='gray', linestyle='--', alpha=0.5, label='Standard dose')\n",
    "axes[1, 1].set_xlabel('Treatment Cycle', fontsize=11)\n",
    "axes[1, 1].set_ylabel('Average Dose (fraction of standard)', fontsize=11)\n",
    "axes[1, 1].set_title('Dosing Trajectories Over Treatment', fontsize=13, fontweight='bold')\n",
    "axes[1, 1].legend(fontsize=10)\n",
    "axes[1, 1].set_ylim(0.5, 1.2)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ],
   "id": "cell_16"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Individual Patient Case Studies"
   ],
   "id": "cell_17"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show detailed dosing trajectory for 3 sample patients\n",
    "fig, axes = plt.subplots(3, 3, figsize=(16, 12))\n",
    "\n",
    "for patient_idx in range(3):\n",
    "    seed = 10000 + patient_idx * 7\n",
    "\n",
    "    for agent_idx, (agent_obj, agent_name, color) in enumerate([\n",
    "        (fixed_agent, 'Fixed', 'red'),\n",
    "        (rule_agent, 'Rule-Based', 'orange'),\n",
    "    ]):\n",
    "        env = ChemoDoseEnv()\n",
    "        state = env.reset(seed=seed)\n",
    "        if hasattr(agent_obj, 'reset'):\n",
    "            agent_obj.reset()\n",
    "\n",
    "        tumors, wbcs, doses = [env.patient.tumor_size], [env.patient.wbc], []\n",
    "\n",
    "        for step in range(8):\n",
    "            action = agent_obj.choose_action(state)\n",
    "            state, reward, done, info = env.step(action)\n",
    "            tumors.append(info['tumor_size'])\n",
    "            wbcs.append(info['wbc'])\n",
    "            doses.append(info['dose'])\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "    # Q-learning agent\n",
    "    env = ChemoDoseEnv()\n",
    "    state = env.reset(seed=seed)\n",
    "    q_agent.epsilon = 0.0\n",
    "\n",
    "    q_tumors, q_wbcs, q_doses = [env.patient.tumor_size], [env.patient.wbc], []\n",
    "\n",
    "    for step in range(8):\n",
    "        action = q_agent.choose_action(state)\n",
    "        state, reward, done, info = env.step(action)\n",
    "        q_tumors.append(info['tumor_size'])\n",
    "        q_wbcs.append(info['wbc'])\n",
    "        q_doses.append(info['dose'])\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "    # Plot tumor trajectory\n",
    "    axes[patient_idx, 0].plot(q_tumors, 'b-o', label='Q-Learning', markersize=4)\n",
    "    axes[patient_idx, 0].set_ylabel('Tumor Size (mm)', fontsize=10)\n",
    "    axes[patient_idx, 0].set_title(f'Patient {patient_idx+1}: Tumor Trajectory', fontsize=11, fontweight='bold')\n",
    "    axes[patient_idx, 0].legend(fontsize=9)\n",
    "    axes[patient_idx, 0].grid(True, alpha=0.3)\n",
    "\n",
    "    # Plot WBC trajectory\n",
    "    axes[patient_idx, 1].plot(q_wbcs, 'b-o', label='Q-Learning', markersize=4)\n",
    "    axes[patient_idx, 1].axhline(y=2000, color='red', linestyle='--', alpha=0.5, label='Grade 3 threshold')\n",
    "    axes[patient_idx, 1].set_ylabel('WBC (cells/uL)', fontsize=10)\n",
    "    axes[patient_idx, 1].set_title(f'Patient {patient_idx+1}: WBC Trajectory', fontsize=11, fontweight='bold')\n",
    "    axes[patient_idx, 1].legend(fontsize=9)\n",
    "    axes[patient_idx, 1].grid(True, alpha=0.3)\n",
    "\n",
    "    # Plot dose decisions\n",
    "    axes[patient_idx, 2].bar(range(1, len(q_doses)+1), q_doses, color='blue', alpha=0.7, label='Q-Learning')\n",
    "    axes[patient_idx, 2].axhline(y=1.0, color='gray', linestyle='--', alpha=0.5)\n",
    "    axes[patient_idx, 2].set_ylabel('Dose Fraction', fontsize=10)\n",
    "    axes[patient_idx, 2].set_title(f'Patient {patient_idx+1}: Dose Decisions', fontsize=11, fontweight='bold')\n",
    "    axes[patient_idx, 2].set_ylim(0.4, 1.2)\n",
    "    axes[patient_idx, 2].legend(fontsize=9)\n",
    "    axes[patient_idx, 2].grid(True, alpha=0.3)\n",
    "\n",
    "for ax in axes[2, :]:\n",
    "    ax.set_xlabel('Treatment Cycle', fontsize=10)\n",
    "\n",
    "plt.suptitle('Q-Learning Agent: Individual Patient Trajectories', fontsize=14, fontweight='bold', y=1.01)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"The Q-learning agent adapts its dosing based on each patient's response.\")\n",
    "print(\"Notice how it reduces dose when toxicity rises and increases when the patient tolerates treatment well.\")"
   ],
   "id": "cell_18"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary and Key Findings"
   ],
   "id": "cell_19"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"=\" * 60)\n",
    "print(\"CASE STUDY SUMMARY\")\n",
    "print(\"=\" * 60)\n",
    "print()\n",
    "print(\"Problem: Adaptive chemotherapy dose optimization\")\n",
    "print(\"Company: NovaCure Therapeutics (simulated)\")\n",
    "print(\"Drug:    NC-4817 for advanced non-small-cell lung cancer\")\n",
    "print()\n",
    "print(\"RL Formulation:\")\n",
    "print(\"  States:  tumor size, WBC, toxicity grade, kidney function, cycle\")\n",
    "print(\"  Actions: 6 dose levels (0.6x to 1.1x standard)\")\n",
    "print(\"  Reward:  tumor shrinkage - toxicity penalty + completion bonus\")\n",
    "print(\"  Method:  Tabular Q-learning (gamma=0.95, epsilon-greedy)\")\n",
    "print()\n",
    "print(\"Key Results:\")\n",
    "print(f\"  Q-Learning avg reward:     {np.mean(q_results['rewards']):>7.2f}\")\n",
    "print(f\"  Fixed dose avg reward:     {np.mean(fixed_results['rewards']):>7.2f}\")\n",
    "print(f\"  Rule-based avg reward:     {np.mean(rule_results['rewards']):>7.2f}\")\n",
    "print()\n",
    "improvement = (np.mean(q_results['rewards']) - np.mean(fixed_results['rewards'])) / abs(np.mean(fixed_results['rewards'])) * 100\n",
    "print(f\"  Improvement over fixed dose: {improvement:+.1f}%\")\n",
    "print()\n",
    "print(\"The RL agent learned to personalize dosing for each patient,\")\n",
    "print(\"balancing tumor control against toxicity risk -- exactly the kind\")\n",
    "print(\"of sequential decision problem that RL was designed to solve.\")"
   ],
   "id": "cell_20"
  }
 ]
}