{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1000000",
   "metadata": {},
   "source": [
    "# Human vs Machine: CalMS21 Task 1 Baseline\n",
    "\n",
    "This notebook applies BANOS to a **human-vs-machine** annotation setting using the\n",
    "CalMS21 Task 1 benchmark dataset.\n",
    "\n",
    "**What we do:**\n",
    "1. Load the CalMS21 Task 1 baseline model (Conv-1D trained on pose keypoints)\n",
    "   and run inference on the 19 test recordings\n",
    "2. Reproduce CalMS21's published pooled macro-F1, confirming we replicate\n",
    "   their evaluation pipeline exactly\n",
    "3. Show our preferred per-recording frame F1 + BANOS metrics side-by-side,\n",
    "   explaining deliberate differences from CalMS21's approach\n",
    "\n",
    "**Python-only** (no MATLAB equivalent) — requires `tensorflow`.\n",
    "\n",
    "> For BANOS fundamentals, see [`tutorial_python.ipynb`](tutorial_python.ipynb).\n",
    "\n",
    "**Install dependencies:**\n",
    "```bash\n",
    "# Option A — via the calms21 extra (recommended)\n",
    "pip install \"banos[calms21]\"\n",
    "\n",
    "# Option B — manually\n",
    "pip install tensorflow scikit-learn tqdm\n",
    "```\n",
    "\n",
    "**Launch with uv (dev environment):**\n",
    "```bash\n",
    "uv run --extra calms21 --extra notebook jupyter notebook demo/tutorial_calms21_task1.ipynb\n",
    "```\n",
    "\n",
    "**Data:** The CalMS21 test JSON (~600 MB) is downloaded automatically on first run\n",
    "and cached in `.cache/calms21/` (gitignored). The model weights are included in\n",
    "the repo at `demo/models/task1_seed_42_model.h5`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000001",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 1 — Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a1000002",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow 2.21.0\n",
      "BANOS 0.2.0\n",
      "Model: C:\\Users\\GA401\\Documents\\Pro\\Code\\GithubPerso\\Tl_BANOS_dev\\demo\\models\\task1_seed_42_model.h5  (959 KB)\n",
      "Cache: C:\\Users\\GA401\\Documents\\Pro\\Code\\GithubPerso\\Tl_BANOS_dev\\.cache\\calms21\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import zipfile\n",
    "from pathlib import Path\n",
    "import math\n",
    "import urllib.request\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "try:\n",
    "    import tensorflow as tf\n",
    "    print(f\"TensorFlow {tf.__version__}\")\n",
    "except ImportError:\n",
    "    print(\"TensorFlow not found. Install with: pip install banos[calms21]\")\n",
    "    sys.exit(1)\n",
    "\n",
    "import banos\n",
    "from sklearn.metrics import f1_score\n",
    "\n",
    "print(f\"BANOS {banos.__version__}\")\n",
    "\n",
    "\n",
    "def _find_demo_dir():\n",
    "    \"\"\"Locate the demo/ directory regardless of where Jupyter was launched.\"\"\"\n",
    "    try:\n",
    "        nb_file = Path(__file__).resolve()\n",
    "        if nb_file.parent.name == \"demo\":\n",
    "            return nb_file.parent\n",
    "    except NameError:\n",
    "        pass\n",
    "    cwd = Path(\".\").resolve()\n",
    "    for candidate in [cwd, cwd.parent]:\n",
    "        if (candidate / \"models\" / \"task1_seed_42_model.h5\").exists():\n",
    "            return candidate\n",
    "    return cwd\n",
    "\n",
    "\n",
    "DEMO_DIR = _find_demo_dir()\n",
    "\n",
    "MODEL_PATH = DEMO_DIR / \"models\" / \"task1_seed_42_model.h5\"\n",
    "\n",
    "# CalMS21 Task 1 classification zip (~457 MB) — contains the test JSON\n",
    "TASK1_ZIP_URL  = \"https://data.caltech.edu/records/s0vdx-0k302/files/task1_classic_classification.zip?download=1\"\n",
    "CACHE_DIR      = DEMO_DIR.parent / \".cache\" / \"calms21\"\n",
    "CACHE_DIR.mkdir(parents=True, exist_ok=True)\n",
    "TASK1_ZIP_PATH = CACHE_DIR / \"task1_classic_classification.zip\"\n",
    "TEST_JSON_PATH = CACHE_DIR / \"calms21_task1_test.json\"\n",
    "\n",
    "BEHAVIORS = [\"attack\", \"investigation\", \"mount\"]\n",
    "VOCAB     = {\"attack\": 0, \"investigation\": 1, \"mount\": 2, \"other\": 3}\n",
    "\n",
    "assert MODEL_PATH.exists(), (\n",
    "    f\"Model not found at {MODEL_PATH}\\n\"\n",
    "    f\"Make sure you cloned the full repo (demo/models/task1_seed_42_model.h5 must exist).\"\n",
    ")\n",
    "print(f\"Model: {MODEL_PATH}  ({MODEL_PATH.stat().st_size / 1e3:.0f} KB)\")\n",
    "print(f\"Cache: {CACHE_DIR}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000003",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 2 — Download CalMS21 Test Data\n",
    "\n",
    "The CalMS21 Task 1 classification zip (~457 MB) is downloaded once from CaltechDATA,\n",
    "then the test JSON is extracted from it. Both are cached in `.cache/calms21/` (gitignored).\n",
    "\n",
    "If the download fails, download `task1_classic_classification.zip` manually from\n",
    "https://data.caltech.edu/records/s0vdx-0k302 and place it at the path printed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a1000004",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CalMS21 task1 zip: downloading from https://data.caltech.edu/records/s0vdx-0k302/files/task1_classic_classification.zip?download=1 ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "CalMS21 task1 zip: 457MB [00:41, 11.1MB/s]                                                                    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CalMS21 task1 zip: done (457.4 MB)\n",
      "Extracting calms21_task1_test.json from zip ...\n",
      "Test JSON: extracted (632.2 MB)\n"
     ]
    }
   ],
   "source": [
    "def download_with_progress(url, dest_path, label):\n",
    "    \"\"\"Download a file once; skip if already present.\"\"\"\n",
    "    dest_path = Path(dest_path)\n",
    "    if dest_path.exists():\n",
    "        print(f\"{label}: already cached ({dest_path.stat().st_size / 1e6:.1f} MB)\")\n",
    "        return\n",
    "    print(f\"{label}: downloading from {url} ...\")\n",
    "    try:\n",
    "        from tqdm import tqdm\n",
    "\n",
    "        class _TqdmHook(tqdm):\n",
    "            def update_to(self, b=1, bsize=1, tsize=None):\n",
    "                if tsize is not None:\n",
    "                    self.total = tsize\n",
    "                self.update(b * bsize - self.n)\n",
    "\n",
    "        with _TqdmHook(unit=\"B\", unit_scale=True, miniters=1, desc=label) as t:\n",
    "            urllib.request.urlretrieve(url, dest_path, reporthook=t.update_to)\n",
    "    except ImportError:\n",
    "        def _hook(count, block_size, total_size):\n",
    "            if total_size > 0 and count % 500 == 0:\n",
    "                pct = min(100.0, 100 * count * block_size / total_size)\n",
    "                print(f\"  {pct:.0f}%\", end=\"\\r\")\n",
    "        urllib.request.urlretrieve(url, dest_path, reporthook=_hook)\n",
    "        print()\n",
    "    print(f\"{label}: done ({dest_path.stat().st_size / 1e6:.1f} MB)\")\n",
    "\n",
    "\n",
    "def extract_test_json(zip_path, dest_path):\n",
    "    \"\"\"Extract calms21_task1_test.json from the zip into dest_path.\"\"\"\n",
    "    if dest_path.exists():\n",
    "        print(f\"Test JSON: already extracted ({dest_path.stat().st_size / 1e6:.1f} MB)\")\n",
    "        return\n",
    "    print(\"Extracting calms21_task1_test.json from zip ...\")\n",
    "    with zipfile.ZipFile(zip_path) as zf:\n",
    "        # Find the test JSON entry (may be in a subdirectory inside the zip)\n",
    "        matches = [n for n in zf.namelist() if n.endswith(\"calms21_task1_test.json\")]\n",
    "        if not matches:\n",
    "            raise FileNotFoundError(\n",
    "                \"calms21_task1_test.json not found in zip. \"\n",
    "                f\"Zip contains: {zf.namelist()[:10]} ...\"\n",
    "            )\n",
    "        entry = matches[0]\n",
    "        with zf.open(entry) as src, open(dest_path, \"wb\") as dst:\n",
    "            import shutil\n",
    "            shutil.copyfileobj(src, dst)\n",
    "    print(f\"Test JSON: extracted ({dest_path.stat().st_size / 1e6:.1f} MB)\")\n",
    "\n",
    "\n",
    "download_with_progress(TASK1_ZIP_URL, TASK1_ZIP_PATH, \"CalMS21 task1 zip\")\n",
    "extract_test_json(TASK1_ZIP_PATH, TEST_JSON_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000005",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 3 — Parse Test JSON"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a1000006",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading test JSON (this may take a moment for a 600 MB file)...\n",
      "Number of test recordings: 19\n",
      "Vocabulary confirmed: {'attack': 0, 'investigation': 1, 'mount': 2, 'other': 3}\n",
      "Total frames: 262,107\n",
      "\n",
      "  mouse071_task1_annotator1: 23810 frames | investigation=4679  mount=2894  other=16237\n",
      "  mouse072_task1_annotator1: 18164 frames | investigation=4298  mount=6409  other=7457\n",
      "  mouse073_task1_annotator1: 19156 frames | investigation=6867  mount=5050  other=7239\n",
      "  ...\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "print(\"Loading test JSON (this may take a moment for a 600 MB file)...\")\n",
    "with open(TEST_JSON_PATH) as f:\n",
    "    raw = json.load(f)\n",
    "\n",
    "# Structure: { 'annotator-id_0': { 'task1/test/<name>': { keypoints, annotations, ... } } }\n",
    "seqs       = raw[\"annotator-id_0\"]\n",
    "recordings = sorted(seqs.keys())\n",
    "\n",
    "print(f\"Number of test recordings: {len(recordings)}\")\n",
    "\n",
    "vocab_check = seqs[recordings[0]][\"metadata\"][\"vocab\"]\n",
    "assert vocab_check == VOCAB, f\"Unexpected vocab: {vocab_check}\"\n",
    "print(f\"Vocabulary confirmed: {VOCAB}\")\n",
    "\n",
    "dataset = {\n",
    "    rec: {\n",
    "        \"keypoints\": np.array(seqs[rec][\"keypoints\"]),   # (n_frames, 2, 2, 7)\n",
    "        \"gt\":        np.array(seqs[rec][\"annotations\"]),  # (n_frames,) 4-class int\n",
    "    }\n",
    "    for rec in recordings\n",
    "}\n",
    "\n",
    "total_frames = sum(v[\"gt\"].shape[0] for v in dataset.values())\n",
    "print(f\"Total frames: {total_frames:,}\")\n",
    "print()\n",
    "for rec in recordings[:3]:\n",
    "    n = dataset[rec][\"gt\"].shape[0]\n",
    "    labels, counts = np.unique(dataset[rec][\"gt\"], return_counts=True)\n",
    "    beh_str = \"  \".join(f\"{list(VOCAB)[l]}={int(c)}\" for l, c in zip(labels, counts))\n",
    "    print(f\"  {rec.split('/')[-1]}: {n} frames | {beh_str}\")\n",
    "print(\"  ...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000007",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 4 — Preprocess, Build Windows, and Run Inference\n",
    "\n",
    "The CalMS21 baseline preprocessing pipeline:\n",
    "1. **Transpose** `(n, 2, 2, 7)` → `(n, 2, 7, 2)` so the last axis is `(x, y)`\n",
    "2. **Normalize** `x / 1024`, `y / 570` (original image resolution is 1024 × 570 px)\n",
    "3. **Flatten** → `(n, 28)` — 2 mice × 7 body parts × 2 coords\n",
    "\n",
    "Each frame then becomes a **201-frame window** (100 past × step 2, centre, 100 future × step 2).\n",
    "Boundaries are zero-padded, matching the original training generator.\n",
    "\n",
    "Windows are built one recording at a time to keep peak RAM usage ~500 MB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a1000008",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model...\n",
      "Model loaded. Running inference on 19 test recordings...\n",
      "\n",
      "WARNING:tensorflow:TensorFlow GPU support is not available on native Windows for TensorFlow >= 2.11. Even if CUDA/cuDNN are installed, GPU will not be used. Please use WSL2 or the TensorFlow-DirectML plugin.\n",
      "  mouse071_task1_annotator1: 23,810 frames\n",
      "  mouse072_task1_annotator1: 18,164 frames\n",
      "  mouse073_task1_annotator1: 19,156 frames\n",
      "  mouse074_task1_annotator1: 18,825 frames\n",
      "  mouse075_task1_annotator1: 18,015 frames\n",
      "  mouse076_task1_annotator1: 19,710 frames\n",
      "  mouse077_task1_annotator1: 18,907 frames\n",
      "  mouse078_task1_annotator1: 19,473 frames\n",
      "  mouse079_task1_annotator1: 17,483 frames\n",
      "  mouse080_task1_annotator1: 6,179 frames\n",
      "  mouse081_task1_annotator1: 4,932 frames\n",
      "  mouse082_task1_annotator1: 7,115 frames\n",
      "  mouse083_task1_annotator1: 4,881 frames\n",
      "  mouse084_task1_annotator1: 5,626 frames\n",
      "  mouse085_task1_annotator1: 19,200 frames\n",
      "  mouse086_task1_annotator1: 8,617 frames\n",
      "  mouse087_task1_annotator1: 11,915 frames\n",
      "  mouse088_task1_annotator1: 8,825 frames\n",
      "  mouse089_task1_annotator1: 11,274 frames\n",
      "\n",
      "Inference complete.\n"
     ]
    }
   ],
   "source": [
    "def preprocess_keypoints(kp):\n",
    "    \"\"\"Normalize and flatten keypoints to shape (n_frames, 28).\"\"\"\n",
    "    kp = kp.transpose(0, 1, 3, 2).astype(np.float32)  # (n, 2, 7, 2)\n",
    "    kp[..., 0] /= 1024.0\n",
    "    kp[..., 1] /= 570.0\n",
    "    return kp.reshape(kp.shape[0], -1)                 # (n, 28)\n",
    "\n",
    "\n",
    "def make_windows(kp_norm, past=100, future=100, gap=2):\n",
    "    \"\"\"\n",
    "    Build sliding temporal windows matching the CalMS21 baseline data generator.\n",
    "\n",
    "    kp_norm : ndarray (n_frames, 28)\n",
    "    Returns : ndarray (n_frames, past + 1 + future, 28)  — i.e. (n, 201, 28)\n",
    "    \"\"\"\n",
    "    n, feat_dim = kp_norm.shape\n",
    "    left_pad  = past  * gap   # 200 zero frames on the left\n",
    "    right_pad = future * gap  # 200 zero frames on the right\n",
    "    padded = np.pad(kp_norm, ((left_pad, right_pad), (0, 0)))  # (n + 400, 28)\n",
    "\n",
    "    # Vectorised index construction — no Python loop over frames\n",
    "    offsets     = np.arange(0, (past + 1 + future) * gap, gap)  # (201,)\n",
    "    all_indices = np.arange(n)[:, None] + offsets                # (n, 201)\n",
    "    return padded[all_indices]  # (n, 201, 28)\n",
    "\n",
    "\n",
    "# compile=False: avoids needing tensorflow-addons (used only for training metrics)\n",
    "print(\"Loading model...\")\n",
    "model = tf.keras.models.load_model(MODEL_PATH, compile=False)\n",
    "print(\"Model loaded. Running inference on 19 test recordings...\")\n",
    "print()\n",
    "\n",
    "for rec in recordings:\n",
    "    kp_norm = preprocess_keypoints(dataset[rec][\"keypoints\"])   # (n, 28)\n",
    "    X       = make_windows(kp_norm)                              # (n, 201, 28)\n",
    "    probs   = model.predict(X, batch_size=512, verbose=0)        # (n, 4)\n",
    "    dataset[rec][\"pred\"] = np.argmax(probs, axis=-1)             # (n,)\n",
    "    print(f\"  {rec.split('/')[-1]}: {len(probs):,} frames\")\n",
    "    del X, probs  # free window array before next recording\n",
    "\n",
    "print(\"\\nInference complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000009",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 5 — Convert to Binary DataFrames\n",
    "\n",
    "BANOS expects one binary column per behavior. We convert 4-class integer labels\n",
    "to three binary columns (attack / investigation / mount)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a1000010",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example — mouse071_task1_annotator1\n",
      "  GT bout counts:   {'attack': 0, 'investigation': 79, 'mount': 26}\n",
      "  Pred bout counts: {'attack': 1, 'investigation': 76, 'mount': 29}\n"
     ]
    }
   ],
   "source": [
    "banos_data = {}\n",
    "for rec in recordings:\n",
    "    gt_4   = dataset[rec][\"gt\"]\n",
    "    pred_4 = dataset[rec][\"pred\"]\n",
    "    gt_df   = pd.DataFrame({b: (gt_4   == VOCAB[b]).astype(int) for b in BEHAVIORS})\n",
    "    pred_df = pd.DataFrame({b: (pred_4 == VOCAB[b]).astype(int) for b in BEHAVIORS})\n",
    "    banos_data[rec] = (pred_df, gt_df)\n",
    "\n",
    "rec0 = recordings[0]\n",
    "p0, g0 = banos_data[rec0]\n",
    "print(f\"Example — {rec0.split('/')[-1]}\")\n",
    "print(f\"  GT bout counts:   { {b: int((g0[b].diff().fillna(g0[b]) == 1).sum()) for b in BEHAVIORS} }\")\n",
    "print(f\"  Pred bout counts: { {b: int((p0[b].diff().fillna(p0[b]) == 1).sum()) for b in BEHAVIORS} }\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000011",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 6 — CalMS21 Official F1 (Pooled Method)\n",
    "\n",
    "CalMS21 evaluates by **pooling all frames** from all recordings, then computing\n",
    "per-class F1 with `sklearn.metrics.f1_score(average=None)`.\n",
    "Macro-F1 = mean of attack / investigation / mount (indices 0, 1, 2);\n",
    "\"other\" (index 3) is excluded.\n",
    "\n",
    "We replicate this exactly to confirm our model and pipeline are correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a1000012",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CalMS21 Official F1 (pooled frames, sklearn, macro excluding 'other'):\n",
      "  attack           F1 = 0.6318\n",
      "  investigation    F1 = 0.8057\n",
      "  mount            F1 = 0.9000\n",
      "  other            F1 = 0.9456  (excluded)\n",
      "\n",
      "  Macro F1 (excl. other): 0.7792\n"
     ]
    }
   ],
   "source": [
    "all_gt   = np.concatenate([dataset[r][\"gt\"]   for r in recordings])\n",
    "all_pred = np.concatenate([dataset[r][\"pred\"] for r in recordings])\n",
    "\n",
    "f1_per_class = f1_score(all_gt, all_pred, average=None,\n",
    "                        labels=[0, 1, 2, 3], zero_division=0)\n",
    "\n",
    "calms21_macro_f1 = f1_per_class[:3].mean()   # exclude \"other\"\n",
    "\n",
    "print(\"CalMS21 Official F1 (pooled frames, sklearn, macro excluding 'other'):\")\n",
    "for beh, f1v in zip(BEHAVIORS, f1_per_class[:3]):\n",
    "    print(f\"  {beh:<16} F1 = {f1v:.4f}\")\n",
    "print(f\"  {'other':<16} F1 = {f1_per_class[3]:.4f}  (excluded)\")\n",
    "print(f\"\\n  Macro F1 (excl. other): {calms21_macro_f1:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000013",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 7 — Our Per-Recording Frame F1\n",
    "\n",
    "We compute frame F1 **per recording, per behavior**, macro-average across\n",
    "behaviors, then mean across recordings.\n",
    "\n",
    "**Two deliberate differences from CalMS21:**\n",
    "1. **Per-recording aggregation** — treats each recording equally (CalMS21 weights by length)\n",
    "2. **Explicit absent-behavior scoring**  \n",
    "   - GT = 0 bouts AND machine = 0 bouts → **1.0** (correct absence)  \n",
    "   - GT = 0 bouts BUT machine predicts bouts → **0.0** (false detection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a1000014",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Per-recording Frame F1 (mean across recordings):\n",
      "  attack           mean F1 = 0.6325\n",
      "  investigation    mean F1 = 0.7441\n",
      "  mount            mean F1 = 0.8085\n",
      "\n",
      "  Overall (macro of means): 0.7267\n"
     ]
    }
   ],
   "source": [
    "def frame_f1(pred_col, gt_col):\n",
    "    \"\"\"Frame-level F1 with explicit absent-behavior scoring.\"\"\"\n",
    "    gt_col, pred_col = np.asarray(gt_col), np.asarray(pred_col)\n",
    "    if gt_col.sum() == 0 and pred_col.sum() == 0:\n",
    "        return 1.0\n",
    "    if gt_col.sum() == 0 and pred_col.sum() > 0:\n",
    "        return 0.0\n",
    "    tp = int(((pred_col == 1) & (gt_col == 1)).sum())\n",
    "    fp = int(((pred_col == 1) & (gt_col == 0)).sum())\n",
    "    fn = int(((pred_col == 0) & (gt_col == 1)).sum())\n",
    "    p  = tp / (tp + fp) if tp + fp > 0 else float(\"nan\")\n",
    "    r  = tp / (tp + fn) if tp + fn > 0 else float(\"nan\")\n",
    "    return 2 * p * r / (p + r) if (p + r) > 0 else float(\"nan\")\n",
    "\n",
    "\n",
    "per_rec_f1 = {}\n",
    "for rec in recordings:\n",
    "    p_df, g_df = banos_data[rec]\n",
    "    beh_f1s = {b: frame_f1(p_df[b].values, g_df[b].values) for b in BEHAVIORS}\n",
    "    per_rec_f1[rec] = {**beh_f1s, \"macro\": np.nanmean(list(beh_f1s.values()))}\n",
    "\n",
    "our_frame_f1 = np.nanmean([v[\"macro\"] for v in per_rec_f1.values()])\n",
    "\n",
    "print(\"Per-recording Frame F1 (mean across recordings):\")\n",
    "for b in BEHAVIORS:\n",
    "    beh_mean = np.nanmean([per_rec_f1[r][b] for r in recordings])\n",
    "    print(f\"  {b:<16} mean F1 = {beh_mean:.4f}\")\n",
    "print(f\"\\n  Overall (macro of means): {our_frame_f1:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000015",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 8 — BANOS Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a1000016",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BANOS per-behavior metrics (nanmean across 19 recordings):\n",
      "               DA (F1)      SO      TP      IC  Precision  Recall\n",
      "attack          0.5745  0.5272  0.4800  0.7287     0.5548  0.6076\n",
      "investigation   0.6476  0.4556  0.0651  0.9734     0.6166  0.7061\n",
      "mount           0.7006  0.6357  0.2586  0.9374     0.6801  0.7264\n",
      "\n",
      "BANOS overall metrics (nanmean across behaviors and recordings):\n",
      "  precision    0.6172\n",
      "  recall       0.6801\n",
      "  f1_score     0.6398\n",
      "  so           0.5378\n",
      "  tp           0.2681\n",
      "  ic           0.8798\n"
     ]
    }
   ],
   "source": [
    "group_metrics, overall_metrics = banos.score(banos_data)\n",
    "\n",
    "print(\"BANOS per-behavior metrics (nanmean across 19 recordings):\")\n",
    "group_df = pd.DataFrame(group_metrics).T[\n",
    "    [\"f1_score\", \"so\", \"tp\", \"ic\", \"precision\", \"recall\"]\n",
    "]\n",
    "group_df.columns = [\"DA (F1)\", \"SO\", \"TP\", \"IC\", \"Precision\", \"Recall\"]\n",
    "print(group_df.round(4).to_string())\n",
    "\n",
    "print(\"\\nBANOS overall metrics (nanmean across behaviors and recordings):\")\n",
    "for metric, value in overall_metrics.items():\n",
    "    v = f\"{value:.4f}\" if (value is not None and not math.isnan(value)) else \"NaN\"\n",
    "    print(f\"  {metric:<12} {v}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000017",
   "metadata": {},
   "source": [
    "---\n",
    "## Section 9 — Side-by-Side Comparison\n",
    "\n",
    "### Why the numbers differ\n",
    "\n",
    "| Method | Aggregation | Absent behavior |\n",
    "|--------|-------------|------------------|\n",
    "| **CalMS21 macro-F1** | Pool all frames globally | Implicit (rare behaviors diluted into pool) |\n",
    "| **Our frame F1** | Per recording → mean | Explicit: correct absence = 1.0, false detection = 0.0 |\n",
    "| **BANOS DA** | Per recording → nanmean | Same as our frame F1 |\n",
    "\n",
    "BANOS additionally reveals **how** the model errs:\n",
    "- Low **SO** → predicted bouts have poor temporal overlap with GT bouts\n",
    "- Low **TP** → bout boundaries (start/end frames) are imprecise\n",
    "- High **IC** → predictions are stable once they fire (no within-bout flickering)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a1000018",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metric                     Value  What it captures\n",
      "---------------------------------------------------------------------------\n",
      "CalMS21 F1 (pooled)       0.7792  Global frame overlap; matches leaderboard method\n",
      "Our Frame F1              0.7267  Per-recording avg; explicit absent-behavior scoring\n",
      "BANOS DA (F1)             0.6398  Bout-level detection accuracy\n",
      "BANOS SO                  0.5378  Temporal overlap quality of matched bouts\n",
      "BANOS TP                  0.2681  Boundary precision (start/end frame accuracy)\n",
      "BANOS IC                  0.8798  Label stability within detected bouts\n"
     ]
    }
   ],
   "source": [
    "banos_da = overall_metrics[\"f1_score\"]\n",
    "banos_so = overall_metrics[\"so\"]\n",
    "banos_tp = overall_metrics[\"tp\"]\n",
    "banos_ic = overall_metrics[\"ic\"]\n",
    "\n",
    "\n",
    "def _fmt(v):\n",
    "    return f\"{v:.4f}\" if (v is not None and not math.isnan(v)) else \"NaN\"\n",
    "\n",
    "\n",
    "comparison = [\n",
    "    (\"CalMS21 F1 (pooled)\",  calms21_macro_f1,\n",
    "     \"Global frame overlap; matches leaderboard method\"),\n",
    "    (\"Our Frame F1\",         our_frame_f1,\n",
    "     \"Per-recording avg; explicit absent-behavior scoring\"),\n",
    "    (\"BANOS DA (F1)\",        banos_da,\n",
    "     \"Bout-level detection accuracy\"),\n",
    "    (\"BANOS SO\",             banos_so,\n",
    "     \"Temporal overlap quality of matched bouts\"),\n",
    "    (\"BANOS TP\",             banos_tp,\n",
    "     \"Boundary precision (start/end frame accuracy)\"),\n",
    "    (\"BANOS IC\",             banos_ic,\n",
    "     \"Label stability within detected bouts\"),\n",
    "]\n",
    "\n",
    "print(f\"{'Metric':<24} {'Value':>7}  What it captures\")\n",
    "print(\"-\" * 75)\n",
    "for name, val, desc in comparison:\n",
    "    print(f\"{name:<24} {_fmt(val):>7}  {desc}\")"
   ]
  }
 ],
 "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
