{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# MLP Training (ClearML-compatible)\n",
        "\n",
        "PyTorch MLP for focus classification.\n",
        "- Single CFG dict (ClearML `task.connect(CFG)`)\n",
        "- 70/15/15 stratified random split\n",
        "- Per-epoch train/val loss + accuracy table\n",
        "- Test evaluation: accuracy, F1, ROC-AUC\n",
        "- ClearML scalar logging (opt-in)\n",
        "- LOPO comparison at the end"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Imports and CFG"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "import sys\n",
        "import random\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.metrics import (\n",
        "    accuracy_score, f1_score, roc_auc_score,\n",
        "    classification_report, confusion_matrix, ConfusionMatrixDisplay,\n",
        ")\n",
        "import matplotlib.pyplot as plt\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "# Add project root to sys.path\n",
        "_cwd = os.getcwd()\n",
        "PROJECT_ROOT = _cwd if os.path.isdir(os.path.join(_cwd, \"models\")) else os.path.abspath(os.path.join(_cwd, \"..\"))\n",
        "if PROJECT_ROOT not in sys.path:\n",
        "    sys.path.insert(0, PROJECT_ROOT)\n",
        "\n",
        "from data_preparation.prepare_dataset import load_per_person, SELECTED_FEATURES, _split_and_scale\n",
        "\n",
        "CFG = {\n",
        "    \"model_name\": \"face_orientation\",\n",
        "    \"seed\": 42,\n",
        "    \"split_ratios\": (0.7, 0.15, 0.15),\n",
        "    \"scale\": True,\n",
        "    \"batch_size\": 32,\n",
        "    \"epochs\": 30,\n",
        "    \"lr\": 1e-3,\n",
        "    \"hidden_sizes\": [64, 32],\n",
        "    \"checkpoints_dir\": os.path.join(PROJECT_ROOT, \"checkpoints\"),\n",
        "    \"logs_dir\": os.path.join(PROJECT_ROOT, \"evaluation\", \"logs\"),\n",
        "}\n",
        "\n",
        "print(f\"Project root: {PROJECT_ROOT}\")\n",
        "print(f\"Device: {'cuda' if torch.cuda.is_available() else 'cpu'}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. ClearML (optional)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "USE_CLEARML = False  # set True when ClearML credentials are configured\n",
        "task = None\n",
        "\n",
        "if USE_CLEARML:\n",
        "    from clearml import Task\n",
        "    task = Task.init(\n",
        "        project_name=\"FocusGuards Large Group Project\",\n",
        "        task_name=\"MLP Model Training\",\n",
        "        tags=[\"training\", \"mlp\"]\n",
        "    )\n",
        "    task.connect(CFG)\n",
        "    print(\"[ClearML] Connected\")\n",
        "else:\n",
        "    print(\"[ClearML] Disabled (set USE_CLEARML = True to enable)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 3. Load data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def set_seed(seed):\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",
        "set_seed(CFG[\"seed\"])\n",
        "\n",
        "by_person, X_all, y_all = load_per_person(CFG[\"model_name\"])\n",
        "person_names = sorted(by_person.keys())\n",
        "num_features = X_all.shape[1]\n",
        "num_classes = int(y_all.max()) + 1\n",
        "print(f\"\\nPersons: {person_names}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Random split (70/15/15) and scaling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "splits, scaler = _split_and_scale(X_all, y_all, CFG[\"split_ratios\"], CFG[\"seed\"], CFG[\"scale\"])\n",
        "X_train, y_train = splits[\"X_train\"], splits[\"y_train\"]\n",
        "X_val, y_val = splits[\"X_val\"], splits[\"y_val\"]\n",
        "X_test, y_test = splits[\"X_test\"], splits[\"y_test\"]\n",
        "\n",
        "print(f\"Features: {num_features}, Classes: {num_classes}\")\n",
        "\n",
        "def make_loader(X, y, batch_size, shuffle=False):\n",
        "    ds = TensorDataset(torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.long))\n",
        "    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle)\n",
        "\n",
        "train_loader = make_loader(X_train, y_train, CFG[\"batch_size\"], shuffle=True)\n",
        "val_loader = make_loader(X_val, y_val, CFG[\"batch_size\"])\n",
        "test_loader = make_loader(X_test, y_test, CFG[\"batch_size\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 5. Model definition"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class MLP(nn.Module):\n",
        "    def __init__(self, in_features, hidden_sizes, num_classes):\n",
        "        super().__init__()\n",
        "        layers = []\n",
        "        prev = in_features\n",
        "        for h in hidden_sizes:\n",
        "            layers += [nn.Linear(prev, h), nn.ReLU()]\n",
        "            prev = h\n",
        "        layers.append(nn.Linear(prev, num_classes))\n",
        "        self.network = nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.network(x)\n",
        "\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model = MLP(num_features, CFG[\"hidden_sizes\"], num_classes).to(device)\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(model.parameters(), lr=CFG[\"lr\"])\n",
        "\n",
        "param_count = sum(p.numel() for p in model.parameters())\n",
        "print(f\"Model: MLP {[num_features] + CFG['hidden_sizes'] + [num_classes]}\")\n",
        "print(f\"Parameters: {param_count:,}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Training loop (per-epoch logging)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "history = {\n",
        "    \"model_name\": f\"mlp_{CFG['model_name']}\",\n",
        "    \"param_count\": param_count,\n",
        "    \"epochs\": [],\n",
        "    \"train_loss\": [],\n",
        "    \"train_acc\": [],\n",
        "    \"val_loss\": [],\n",
        "    \"val_acc\": [],\n",
        "}\n",
        "best_val_acc = 0.0\n",
        "best_ckpt_path = os.path.join(CFG[\"checkpoints_dir\"], \"mlp_best.pt\")\n",
        "os.makedirs(CFG[\"checkpoints_dir\"], exist_ok=True)\n",
        "\n",
        "print(f\"{'Epoch':>6} | {'Train Loss':>10} | {'Train Acc':>9} | {'Val Loss':>10} | {'Val Acc':>9}\")\n",
        "print(\"-\" * 60)\n",
        "\n",
        "for epoch in range(1, CFG[\"epochs\"] + 1):\n",
        "    model.train()\n",
        "    t_loss, t_correct, t_total = 0.0, 0, 0\n",
        "    for xb, yb in train_loader:\n",
        "        xb, yb = xb.to(device), yb.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        out = model(xb)\n",
        "        loss = criterion(out, yb)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        t_loss += loss.item() * xb.size(0)\n",
        "        t_correct += (out.argmax(1) == yb).sum().item()\n",
        "        t_total += xb.size(0)\n",
        "    train_loss = t_loss / t_total\n",
        "    train_acc = t_correct / t_total\n",
        "\n",
        "    model.eval()\n",
        "    v_loss, v_correct, v_total = 0.0, 0, 0\n",
        "    with torch.no_grad():\n",
        "        for xb, yb in val_loader:\n",
        "            xb, yb = xb.to(device), yb.to(device)\n",
        "            out = model(xb)\n",
        "            loss = criterion(out, yb)\n",
        "            v_loss += loss.item() * xb.size(0)\n",
        "            v_correct += (out.argmax(1) == yb).sum().item()\n",
        "            v_total += xb.size(0)\n",
        "    val_loss = v_loss / v_total\n",
        "    val_acc = v_correct / v_total\n",
        "\n",
        "    history[\"epochs\"].append(epoch)\n",
        "    history[\"train_loss\"].append(round(train_loss, 4))\n",
        "    history[\"train_acc\"].append(round(train_acc, 4))\n",
        "    history[\"val_loss\"].append(round(val_loss, 4))\n",
        "    history[\"val_acc\"].append(round(val_acc, 4))\n",
        "\n",
        "    if task is not None:\n",
        "        task.logger.report_scalar(\"Loss\", \"Train\", float(train_loss), iteration=epoch)\n",
        "        task.logger.report_scalar(\"Loss\", \"Val\", float(val_loss), iteration=epoch)\n",
        "        task.logger.report_scalar(\"Accuracy\", \"Train\", float(train_acc), iteration=epoch)\n",
        "        task.logger.report_scalar(\"Accuracy\", \"Val\", float(val_acc), iteration=epoch)\n",
        "        task.logger.report_scalar(\"Learning Rate\", \"LR\", float(optimizer.param_groups[0][\"lr\"]), iteration=epoch)\n",
        "        task.logger.flush()\n",
        "\n",
        "    marker = \"\"\n",
        "    if val_acc > best_val_acc:\n",
        "        best_val_acc = val_acc\n",
        "        torch.save(model.state_dict(), best_ckpt_path)\n",
        "        marker = \" *\"\n",
        "\n",
        "    print(f\"{epoch:>6} | {train_loss:>10.4f} | {train_acc:>8.2%} | {val_loss:>10.4f} | {val_acc:>8.2%}{marker}\")\n",
        "\n",
        "print(f\"\\nBest val accuracy: {best_val_acc:.2%}\")\n",
        "print(f\"Checkpoint: {best_ckpt_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 7. Loss and accuracy curves"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))\n",
        "epochs = history[\"epochs\"]\n",
        "ax1.plot(epochs, history[\"train_loss\"], label=\"Train\")\n",
        "ax1.plot(epochs, history[\"val_loss\"], label=\"Val\")\n",
        "ax1.set_xlabel(\"Epoch\"); ax1.set_ylabel(\"Loss\"); ax1.set_title(\"Loss\"); ax1.legend()\n",
        "ax2.plot(epochs, history[\"train_acc\"], label=\"Train\")\n",
        "ax2.plot(epochs, history[\"val_acc\"], label=\"Val\")\n",
        "ax2.set_xlabel(\"Epoch\"); ax2.set_ylabel(\"Accuracy\"); ax2.set_title(\"Accuracy\"); ax2.legend()\n",
        "plt.suptitle(f\"MLP Training — {CFG['model_name']}\")\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 8. Test evaluation (random split)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model.load_state_dict(torch.load(best_ckpt_path, weights_only=True))\n",
        "model.eval()\n",
        "\n",
        "all_preds, all_labels, all_probs = [], [], []\n",
        "test_loss_sum, test_total = 0.0, 0\n",
        "with torch.no_grad():\n",
        "    for xb, yb in test_loader:\n",
        "        xb, yb = xb.to(device), yb.to(device)\n",
        "        out = model(xb)\n",
        "        test_loss_sum += criterion(out, yb).item() * xb.size(0)\n",
        "        test_total += xb.size(0)\n",
        "        probs = torch.softmax(out, dim=1)\n",
        "        all_preds.extend(out.argmax(1).cpu().numpy())\n",
        "        all_labels.extend(yb.cpu().numpy())\n",
        "        all_probs.extend(probs.cpu().numpy())\n",
        "\n",
        "test_loss = test_loss_sum / test_total\n",
        "test_preds = np.array(all_preds)\n",
        "test_labels = np.array(all_labels)\n",
        "test_probs = np.array(all_probs)\n",
        "\n",
        "test_acc = float(accuracy_score(test_labels, test_preds))\n",
        "test_f1 = float(f1_score(test_labels, test_preds, average=\"weighted\"))\n",
        "if num_classes > 2:\n",
        "    test_auc = float(roc_auc_score(test_labels, test_probs, multi_class=\"ovr\", average=\"weighted\"))\n",
        "else:\n",
        "    test_auc = float(roc_auc_score(test_labels, test_probs[:, 1]))\n",
        "\n",
        "print(f\"[TEST] Loss:     {test_loss:.4f}\")\n",
        "print(f\"[TEST] Accuracy: {test_acc:.2%}\")\n",
        "print(f\"[TEST] F1:       {test_f1:.4f}\")\n",
        "print(f\"[TEST] ROC-AUC:  {test_auc:.4f}\")\n",
        "\n",
        "if task is not None:\n",
        "    task.logger.report_single_value(\"test_accuracy\", test_acc)\n",
        "    task.logger.report_single_value(\"test_f1\", test_f1)\n",
        "    task.logger.report_single_value(\"test_auc\", test_auc)\n",
        "    task.logger.flush()\n",
        "\n",
        "print(\"\\nClassification report:\")\n",
        "print(classification_report(test_labels, test_preds, target_names=[\"Unfocused (0)\", \"Focused (1)\"]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 9. Confusion matrix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "fig, ax = plt.subplots(figsize=(5, 4))\n",
        "cm = confusion_matrix(test_labels, test_preds)\n",
        "ConfusionMatrixDisplay(cm, display_labels=[\"Unfocused\", \"Focused\"]).plot(ax=ax, cmap=\"Blues\")\n",
        "ax.set_title(f\"MLP confusion matrix — test acc {test_acc:.2%}\")\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 10. Save checkpoint and JSON log"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "history[\"test_loss\"] = round(test_loss, 4)\n",
        "history[\"test_acc\"] = round(test_acc, 4)\n",
        "history[\"test_f1\"] = round(test_f1, 4)\n",
        "history[\"test_auc\"] = round(test_auc, 4)\n",
        "\n",
        "os.makedirs(CFG[\"logs_dir\"], exist_ok=True)\n",
        "log_path = os.path.join(CFG[\"logs_dir\"], f\"mlp_{CFG['model_name']}_training_log.json\")\n",
        "with open(log_path, \"w\") as f:\n",
        "    json.dump(history, f, indent=2)\n",
        "\n",
        "print(f\"[CKPT] Best model: {best_ckpt_path}\")\n",
        "print(f\"[LOG]  History:    {log_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 11. LOPO comparison (MLP)\n",
        "\n",
        "Train+test with Leave-One-Person-Out so we can compare fairly with XGBoost/RF under LOPO."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def train_mlp_on_splits(X_train, y_train, X_test, y_test, cfg, n_features, n_classes):\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    sc = StandardScaler()\n",
        "    X_tr = sc.fit_transform(X_train)\n",
        "    X_te = sc.transform(X_test)\n",
        "\n",
        "    tr_ds = TensorDataset(torch.tensor(X_tr, dtype=torch.float32), torch.tensor(y_train, dtype=torch.long))\n",
        "    te_ds = TensorDataset(torch.tensor(X_te, dtype=torch.float32), torch.tensor(y_test, dtype=torch.long))\n",
        "    tr_loader = DataLoader(tr_ds, batch_size=cfg[\"batch_size\"], shuffle=True)\n",
        "    te_loader = DataLoader(te_ds, batch_size=cfg[\"batch_size\"])\n",
        "\n",
        "    net = MLP(n_features, cfg[\"hidden_sizes\"], n_classes).to(device)\n",
        "    opt = optim.Adam(net.parameters(), lr=cfg[\"lr\"])\n",
        "    crit = nn.CrossEntropyLoss()\n",
        "\n",
        "    for _ in range(cfg[\"epochs\"]):\n",
        "        net.train()\n",
        "        for xb, yb in tr_loader:\n",
        "            xb, yb = xb.to(device), yb.to(device)\n",
        "            opt.zero_grad()\n",
        "            crit(net(xb), yb).backward()\n",
        "            opt.step()\n",
        "\n",
        "    net.eval()\n",
        "    preds_list, probs_list, labels_list = [], [], []\n",
        "    with torch.no_grad():\n",
        "        for xb, yb in te_loader:\n",
        "            xb = xb.to(device)\n",
        "            out = net(xb)\n",
        "            preds_list.extend(out.argmax(1).cpu().numpy())\n",
        "            probs_list.extend(torch.softmax(out, dim=1).cpu().numpy())\n",
        "            labels_list.extend(yb.numpy())\n",
        "\n",
        "    preds = np.array(preds_list)\n",
        "    probs = np.array(probs_list)\n",
        "    labels = np.array(labels_list)\n",
        "    acc = accuracy_score(labels, preds)\n",
        "    f1 = f1_score(labels, preds, average=\"weighted\")\n",
        "    auc = roc_auc_score(labels, probs[:, 1]) if n_classes == 2 else roc_auc_score(labels, probs, multi_class=\"ovr\", average=\"weighted\")\n",
        "    return {\"accuracy\": acc, \"f1\": f1, \"roc_auc\": auc}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "print(\"MLP LOPO evaluation\")\n",
        "print(\"-\" * 60)\n",
        "\n",
        "lopo_results = []\n",
        "for test_person in person_names:\n",
        "    train_persons = [p for p in person_names if p != test_person]\n",
        "    X_tr = np.concatenate([by_person[p][0] for p in train_persons], axis=0)\n",
        "    y_tr = np.concatenate([by_person[p][1] for p in train_persons], axis=0)\n",
        "    X_te, y_te = by_person[test_person]\n",
        "\n",
        "    set_seed(CFG[\"seed\"])\n",
        "    metrics = train_mlp_on_splits(X_tr, y_tr, X_te, y_te, CFG, num_features, num_classes)\n",
        "    metrics[\"test_person\"] = test_person\n",
        "    metrics[\"n_test\"] = len(y_te)\n",
        "    lopo_results.append(metrics)\n",
        "    print(f\"  test={test_person}: acc={metrics['accuracy']:.2%}  F1={metrics['f1']:.4f}  AUC={metrics['roc_auc']:.4f}  (n={len(y_te)})\")\n",
        "\n",
        "print(\"\\nMLP LOPO summary (mean +/- std):\")\n",
        "for m in [\"accuracy\", \"f1\", \"roc_auc\"]:\n",
        "    vals = [r[m] for r in lopo_results]\n",
        "    print(f\"  {m}: {np.mean(vals):.4f} +/- {np.std(vals):.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 12. Random split vs LOPO summary"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "lopo_acc = np.mean([r[\"accuracy\"] for r in lopo_results])\n",
        "lopo_f1 = np.mean([r[\"f1\"] for r in lopo_results])\n",
        "lopo_auc = np.mean([r[\"roc_auc\"] for r in lopo_results])\n",
        "\n",
        "summary = pd.DataFrame([\n",
        "    {\"Method\": \"Random split (70/15/15)\", \"Accuracy\": f\"{test_acc:.2%}\", \"F1\": f\"{test_f1:.4f}\", \"ROC-AUC\": f\"{test_auc:.4f}\"},\n",
        "    {\"Method\": \"LOPO (mean)\", \"Accuracy\": f\"{lopo_acc:.2%}\", \"F1\": f\"{lopo_f1:.4f}\", \"ROC-AUC\": f\"{lopo_auc:.4f}\"},\n",
        "])\n",
        "display(summary)\n",
        "\n",
        "print(\"\\nCompare these MLP LOPO numbers with XGBoost (from xgboost.ipynb).\")\n",
        "print(\"If XGB LOPO > MLP LOPO, XGB generalises better across unseen persons.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 13. Per-person accuracy bar chart"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "fig, ax = plt.subplots(figsize=(10, 4))\n",
        "names_sorted = [r[\"test_person\"] for r in lopo_results]\n",
        "accs = [r[\"accuracy\"] for r in lopo_results]\n",
        "ax.bar(names_sorted, accs, color=\"steelblue\", edgecolor=\"black\")\n",
        "ax.axhline(y=lopo_acc, color=\"red\", linestyle=\"--\", label=f\"Mean = {lopo_acc:.2%}\")\n",
        "ax.set_ylabel(\"Accuracy\")\n",
        "ax.set_xlabel(\"Left-out person\")\n",
        "ax.set_title(\"MLP LOPO: test accuracy per left-out person\")\n",
        "ax.legend()\n",
        "plt.xticks(rotation=45, ha=\"right\")\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.13.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}