{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "50f51079",
   "metadata": {},
   "source": [
    "# Lesson 29: PyTorch training loop activity\n",
    "\n",
    "In this activity, you will modify the basic PyTorch training loop from the lesson 29 demo to add:\n",
    "\n",
    "1. **Batching** - Process the training data in mini-batches instead of all at once\n",
    "2. **Validation** - Track model performance on a held-out validation set during training\n",
    "\n",
    "## Notebook set-up\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13c76bb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Third party imports\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# Set random seeds for reproducibility\n",
    "torch.manual_seed(315)\n",
    "np.random.seed(315)\n",
    "\n",
    "# Check for GPU availability\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f'Using device: {device}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaa50a53",
   "metadata": {},
   "source": [
    "## 1. Load preprocessed data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9304299e",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.read_pickle('https://gperdrizet.github.io/FSA_devops/assets/data/unit4/preprocessed_housing_data.pkl')\n",
    "\n",
    "training_df = data['training_df']\n",
    "testing_df = data['testing_df']\n",
    "features = data['features']\n",
    "label = data['label']\n",
    "\n",
    "print(f'Training samples: {len(training_df)}')\n",
    "print(f'Testing samples: {len(testing_df)}')\n",
    "print(f'Features: {features}')\n",
    "print(f'Label: {label}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3660a060",
   "metadata": {},
   "source": [
    "## 2. Prepare PyTorch tensors and DataLoaders\n",
    "\n",
    "### Task 1: Add batching and validation split\n",
    "\n",
    "Currently, the code below creates tensors for training and testing. Your task is to:\n",
    "\n",
    "1. **Add imports** for `TensorDataset` and `DataLoader` from `torch.utils.data`\n",
    "2. **Split training data** into train and validation sets (e.g., 80/20 split)\n",
    "3. **Create DataLoaders** for batched training\n",
    "\n",
    "**Hints:**\n",
    "- Use `torch.randperm()` to shuffle indices for the split\n",
    "- Use `TensorDataset(X, y)` to combine feature and label tensors\n",
    "- Use `DataLoader(dataset, batch_size=32, shuffle=True)` to create batches\n",
    "- Create separate DataLoaders for training and validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7bdc99c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert dataframes to PyTorch tensors and move to device\n",
    "X_train = torch.tensor(training_df[features].values, dtype=torch.float32).to(device)\n",
    "y_train = torch.tensor(training_df[label].values, dtype=torch.float32).unsqueeze(1).to(device)\n",
    "X_test = torch.tensor(testing_df[features].values, dtype=torch.float32).to(device)\n",
    "y_test = torch.tensor(testing_df[label].values, dtype=torch.float32).unsqueeze(1).to(device)\n",
    "\n",
    "print(f'X_train shape: {X_train.shape}')\n",
    "print(f'y_train shape: {y_train.shape}')\n",
    "print(f'X_test shape: {X_test.shape}')\n",
    "print(f'y_test shape: {y_test.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "219abc95",
   "metadata": {},
   "source": [
    "## 3. Build model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfb7fbbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Linear(8, 64), # Fully connected layer (similar to tf.keras.layers.Dense)\n",
    "    nn.ReLU(),\n",
    "    nn.Dropout(0.2),\n",
    "    nn.Linear(64, 32),\n",
    "    nn.ReLU(),\n",
    "    nn.Dropout(0.2),\n",
    "    nn.Linear(32, 1)\n",
    ").to(device)\n",
    "\n",
    "# Define loss function and optimizer\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-2)\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5823b5be",
   "metadata": {},
   "source": [
    "## 4. Training function\n",
    "\n",
    "### Task 2: Update training loop for batching and validation\n",
    "\n",
    "The current training loop processes all training data at once. Your task is to modify this function to:\n",
    "\n",
    "1. **Accept DataLoaders** instead of raw tensors\n",
    "2. **Iterate over batches** in an inner loop within each epoch\n",
    "3. **Compute validation metrics** after each training epoch\n",
    "\n",
    "**Hints:**\n",
    "- Change function signature to accept `train_loader` and `val_loader` instead of `X_train`, `y_train`\n",
    "- Add an inner `for X_batch, y_batch in train_loader:` loop\n",
    "- Accumulate loss across batches, then average for reporting\n",
    "- Use `model.eval()` and `torch.no_grad()` for validation\n",
    "- Track `val_loss` and `val_r2` in the history dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d821d6cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(\n",
    "    model: nn.Module,\n",
    "    X_train: torch.Tensor,\n",
    "    y_train: torch.Tensor,\n",
    "    criterion: nn.Module,\n",
    "    optimizer: optim.Optimizer,\n",
    "    epochs: int = 50,\n",
    "    print_every: int = 5\n",
    ") -> dict[str, list[float]]:\n",
    "    '''Basic training loop for PyTorch model.\n",
    "    \n",
    "    TODO: Modify this function to add:\n",
    "    1. Batching - process data in mini-batches\n",
    "    2. Validation - track performance on a validation set\n",
    "    '''\n",
    "    \n",
    "    history = {'loss': [], 'r2': []}\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        # Set model to training mode\n",
    "        model.train()\n",
    "        \n",
    "        # Zero the gradients\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        # Forward pass\n",
    "        predictions = model(X_train)\n",
    "        \n",
    "        # Calculate loss\n",
    "        loss = criterion(predictions, y_train)\n",
    "        \n",
    "        # Backward pass\n",
    "        loss.backward()\n",
    "        \n",
    "        # Update weights\n",
    "        optimizer.step()\n",
    "        \n",
    "        # Calculate R²\n",
    "        with torch.no_grad():\n",
    "            ss_res = torch.sum((y_train - predictions) ** 2)\n",
    "            ss_tot = torch.sum((y_train - torch.mean(y_train)) ** 2)\n",
    "            r2 = 1 - (ss_res / ss_tot)\n",
    "        \n",
    "        # Record metrics\n",
    "        history['loss'].append(loss.item())\n",
    "        history['r2'].append(r2.item())\n",
    "        \n",
    "        # Print progress\n",
    "        if (epoch + 1) % print_every == 0 or epoch == 0:\n",
    "            print(f'Epoch {epoch+1}/{epochs} - loss: {loss.item():.4f} - R²: {r2.item():.4f}')\n",
    "    \n",
    "    print('\\nTraining complete.')\n",
    "    return history"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf715640",
   "metadata": {},
   "source": [
    "## 5. Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dee08328",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = train_model(\n",
    "    model=model,\n",
    "    X_train=X_train,\n",
    "    y_train=y_train,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    epochs=100,\n",
    "    print_every=10\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3b088b5",
   "metadata": {},
   "source": [
    "## 6. Learning curves\n",
    "\n",
    "**Note:** Once you add validation, update these plots to show both training and validation metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abf68740",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "\n",
    "axes[0].set_title('Training loss')\n",
    "axes[0].plot(history['loss'])\n",
    "axes[0].set_xlabel('Epoch')\n",
    "axes[0].set_ylabel('Loss (MSE)')\n",
    "\n",
    "axes[1].set_title('Training R²')\n",
    "axes[1].plot(history['r2'])\n",
    "axes[1].set_xlabel('Epoch')\n",
    "axes[1].set_ylabel('R²')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffcd44a5",
   "metadata": {},
   "source": [
    "## 7. Test set evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b6adfed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set model to evaluation mode\n",
    "model.eval()\n",
    "\n",
    "# Make predictions (no gradient calculation needed)\n",
    "with torch.no_grad():\n",
    "    predictions = model(X_test).cpu().numpy().flatten()\n",
    "\n",
    "# Calculate R²\n",
    "ss_res = np.sum((testing_df[label].values - predictions) ** 2)\n",
    "ss_tot = np.sum((testing_df[label].values - np.mean(testing_df[label].values)) ** 2)\n",
    "rsquared = 1 - (ss_res / ss_tot)\n",
    "\n",
    "print(f'Model R² on test set: {rsquared:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bdf9c3d",
   "metadata": {},
   "source": [
    "## 8. Performance analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04a3a765",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(8, 4))\n",
    "\n",
    "axes[0].set_title('Model predictions')\n",
    "axes[0].scatter(\n",
    "    testing_df[label], predictions,\n",
    "    c='black', s=0.5, alpha=0.5\n",
    ")\n",
    "axes[0].plot(\n",
    "    [testing_df[label].min(), testing_df[label].max()],\n",
    "    [testing_df[label].min(), testing_df[label].max()],\n",
    "    color='red', linestyle='--'\n",
    ")\n",
    "axes[0].set_xlabel('True values (standardized)')\n",
    "axes[0].set_ylabel('Predicted values (standardized)')\n",
    "\n",
    "axes[1].set_title('Residuals vs predicted values')\n",
    "axes[1].scatter(\n",
    "    predictions, testing_df[label] - predictions,\n",
    "    c='black', s=0.5, alpha=0.5\n",
    ")\n",
    "axes[1].axhline(0, color='red', linestyle='--')\n",
    "axes[1].set_xlabel('Predicted values (standardized)')\n",
    "axes[1].set_ylabel('Residuals (standardized)')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}