{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ef44ce1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "import Loaders as L\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from lightning.pytorch import seed_everything\n",
    "import Models as M\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "torch.backends.cuda.matmul.allow_tf32 = True  \n",
    "torch.set_float32_matmul_precision('high')\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "def setup_seed(seed):\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    np.random.seed(seed)\n",
    "    seed_everything(seed, workers=True)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    \n",
    "setup_seed(2023)\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "Mean = torch.tensor([30.38144216, 42.03988769, 97.8896116]).view(1,3,1,1)\n",
    "Std = torch.tensor([40.63141752, 44.26910074, 50.29294373]).view(1,3,1,1)\n",
    "\n",
    "df = pd.read_csv(\"./Dataset_RARP_video/dataset_videos_folds.csv\")\n",
    "items = df.to_dict(orient=\"records\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e23696fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "FOLD = 0\n",
    "\n",
    "train_set = df.loc[df[f\"Fold_{FOLD}\"] == \"train\"].sort_values(by=[\"label\", \"case\"]).to_dict(orient=\"records\")\n",
    "val_set = df.loc[df[f\"Fold_{FOLD}\"] == \"val\"].sort_values(by=[\"label\", \"case\"]).to_dict(orient=\"records\")\n",
    "test_set = df.loc[df[f\"Fold_{FOLD}\"] == \"test\"].sort_values(by=[\"label\", \"case\"]).to_dict(orient=\"records\")\n",
    "\n",
    "ckpt_paths = [\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_0/checkpoints/RARP-epoch=20.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_1/checkpoints/RARP-epoch=32.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_2/checkpoints/RARP-epoch=28.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_3/checkpoints/RARP-epoch=27.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_4/checkpoints/RARP-epoch=30.ckpt\"),\n",
    "]\n",
    "\n",
    "Model = M.RARP_NVB_DINO_MultiTask.load_from_checkpoint(ckpt_paths[FOLD]).to(device)\n",
    "Model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8d269f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(train_set),len(val_set),len(test_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "733fd5f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "traintransformT2 = torch.nn.Sequential(\n",
    "    #transforms.RandomErasing(0.65, value=\"random\"),\n",
    "    transforms.RandomAffine(degrees=(-15, 15), scale=(0.8, 1.2), fill=5),\n",
    "    transforms.RandomApply([transforms.GaussianBlur(5)], 0.5),\n",
    "    transforms.RandomHorizontalFlip(0.3),\n",
    ").to(device)\n",
    "\n",
    "dataset = L.RARP_Video_Dataset(test_set, (224, 224), (139, 0, 360, 360), decode_resize=(640, 360), mean=Mean, std=Std, transform=traintransformT2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc4a5bca",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = L.RARP_Video_Dataset(test_set, (224, 224), (139, 0, 360, 360), decode_resize=(640, 360), mean=Mean, std=Std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fefae5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = DataLoader(\n",
    "    dataset,\n",
    "    batch_size=4,\n",
    "    shuffle=False,\n",
    "    num_workers=0,\n",
    "    pin_memory=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da6c450",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():    \n",
    "    for video, label in tqdm(loader):\n",
    "        video = video.float().to(device)\n",
    "        label = label.float().to(device)\n",
    "        \n",
    "        frames = video.permute(1, 0, 2, 3, 4)\n",
    "        \n",
    "        for img in frames:\n",
    "            pred, _, _ = Model(img)\n",
    "            pred = pred.flatten()\n",
    "            break\n",
    "        break\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb87eb3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "\n",
    "wind_len = 120\n",
    "start = 0\n",
    "end = (start + wind_len)\n",
    "\n",
    "Mean = [0.485, 0.456, 0.406]\n",
    "Std = [0.229, 0.224, 0.225]\n",
    "\n",
    "device = torch.device(f\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "#arr = np.load(Path(\"Dataset_RARP_video/181/clip_181_30fps.npy\"), mmap_mode=\"r\")\n",
    "arr = np.load(Path(\"frames_from_videos_clips_out/181/clip_181_1fps.npy\"), mmap_mode=\"r\")\n",
    "print(arr.shape)\n",
    "clip_np = arr[start:end].copy()\n",
    "\n",
    "testVal_transform = torch.nn.Sequential(\n",
    "    transforms.CenterCrop(300),\n",
    "    transforms.Resize((224, 224), antialias=True, interpolation=torchvision.transforms.InterpolationMode.BICUBIC),\n",
    "    transforms.Normalize(Mean, Std)\n",
    ").to(device)\n",
    "\n",
    "frame_idx = list(range (0, wind_len-1, 8))\n",
    "\n",
    "frames = clip_np[frame_idx]\n",
    "frames = torch.from_numpy(frames).float() / 255.0\n",
    "frames = testVal_transform(frames)\n",
    "frames = frames * torch.tensor(Std).view(3,1,1) + torch.tensor(Mean).view(3,1,1)\n",
    "frames = torch.clamp(frames, 0, 1)\n",
    "print(frames.shape)\n",
    "\n",
    "grid = torchvision.utils.make_grid(frames, nrow=len(frame_idx))\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.imshow(grid.permute(1,2,0).cpu().numpy())\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8fb2c1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(range (start, end-1, 8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26c3548f",
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_idx = [0, 100, 200, 300, 400, 599] \n",
    "\n",
    "frames = video[1][frame_idx].cpu()\n",
    "frames = frames * Std[0] + Mean[0]\n",
    "frames = torch.clamp(frames[:, [2, 1, 0], :, :]/255, 0, 1)\n",
    "print(frames.shape)\n",
    "\n",
    "grid = torchvision.utils.make_grid(frames, nrow=len(frame_idx))\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.imshow(grid.permute(1,2,0).cpu().numpy())\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ffc335",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# --- INPUT ---\n",
    "# x: torch.Tensor [T, F] = [600, 4]\n",
    "# Replace this with your real tensor:\n",
    "# x = your_tensor  # shape [600, 4]\n",
    "x = torch.randn(600, 4)  # demo only\n",
    "fps = 30                 # set your true fps if you want seconds on the x-axis\n",
    "\n",
    "T, F = x.shape\n",
    "time_idx = np.arange(T)\n",
    "time_sec = time_idx / float(fps)  # for readable axis in seconds\n",
    "\n",
    "# Convert to numpy for Matplotlib convenience\n",
    "arr = x.detach().cpu().numpy()\n",
    "\n",
    "# ---------- 1) Time series per feature with rolling mean & band ----------\n",
    "# Rolling window in frames (e.g., 1 sec = fps, or 3 sec = 3*fps)\n",
    "win = max(3, fps)               # change to your preference (e.g., 15, 30, 60)\n",
    "half = win // 2\n",
    "\n",
    "def rolling_mean_std(a, w):\n",
    "    # valid-mode rolling mean/std\n",
    "    csum = np.cumsum(a, axis=0)\n",
    "    csum = np.pad(csum, ((1,0),(0,0)), mode=\"constant\")  # for easy diff\n",
    "    # window sums:\n",
    "    win_sum = csum[w:] - csum[:-w]\n",
    "    mean = win_sum / float(w)\n",
    "    # rolling std:\n",
    "    # compute rolling sum of squares\n",
    "    sq = a**2\n",
    "    csum_sq = np.cumsum(sq, axis=0)\n",
    "    csum_sq = np.pad(csum_sq, ((1,0),(0,0)), mode=\"constant\")\n",
    "    win_sum_sq = csum_sq[w:] - csum_sq[:-w]\n",
    "    var = (win_sum_sq / float(w)) - mean**2\n",
    "    std = np.sqrt(np.maximum(var, 1e-12))\n",
    "    return mean, std\n",
    "\n",
    "roll_mean, roll_std = rolling_mean_std(arr, win)\n",
    "roll_t = time_sec[half:half+len(roll_mean)]  # align indices\n",
    "\n",
    "fig_ts, axes = plt.subplots(F, 1, figsize=(10, 2.5*F), sharex=True)\n",
    "if F == 1:\n",
    "    axes = [axes]\n",
    "\n",
    "for f in range(F):\n",
    "    ax = axes[f]\n",
    "    ax.plot(time_sec, arr[:, f], alpha=0.35, linewidth=1.0, label=f'Feature {f}')\n",
    "    ax.plot(roll_t, roll_mean[:, f], linewidth=2.0, label=f'Rolling mean (w={win})')\n",
    "    ax.fill_between(roll_t,\n",
    "                    roll_mean[:, f] - roll_std[:, f],\n",
    "                    roll_mean[:, f] + roll_std[:, f],\n",
    "                    alpha=0.2, label='±1 std (rolling)')\n",
    "    ax.set_ylabel(f'Feature {f}')\n",
    "    ax.grid(True, linestyle='--', alpha=0.3)\n",
    "axes[-1].set_xlabel('Time (s)')\n",
    "axes[0].legend(loc='upper right')\n",
    "fig_ts.suptitle('Per-feature time series with rolling mean ± std', y=1.02)\n",
    "fig_ts.tight_layout()\n",
    "fig_ts.savefig('output/time_series_per_feature.png', dpi=200)\n",
    "\n",
    "# ---------- 2) Heatmap overview ----------\n",
    "fig_hm, ax = plt.subplots(figsize=(10, 2.8))\n",
    "# Normalize features to [0,1] for better visual comparison (optional)\n",
    "arr_min = arr.min(axis=0, keepdims=True)\n",
    "arr_max = arr.max(axis=0, keepdims=True)\n",
    "arr_norm = (arr - arr_min) / (arr_max - arr_min + 1e-12)\n",
    "\n",
    "im = ax.imshow(arr_norm.T, aspect='auto', interpolation='nearest',\n",
    "               extent=[time_sec[0], time_sec[-1], F-0.5, -0.5])\n",
    "ax.set_yticks(np.arange(F))\n",
    "ax.set_yticklabels([f'Feat {f}' for f in range(F)])\n",
    "ax.set_xlabel('Time (s)')\n",
    "ax.set_title('Heatmap (per-feature normalized to [0,1])')\n",
    "fig_hm.colorbar(im, ax=ax, fraction=0.025, pad=0.02)\n",
    "fig_hm.tight_layout()\n",
    "fig_hm.savefig('output/heatmap_overview.png', dpi=200)\n",
    "\n",
    "# ---------- 3) Distribution (boxplot per feature) ----------\n",
    "fig_box, ax = plt.subplots(figsize=(7, 3.5))\n",
    "ax.boxplot([arr[:, f] for f in range(F)], showmeans=True)\n",
    "ax.set_xticklabels([f'Feat {f}' for f in range(F)])\n",
    "ax.set_ylabel('Value')\n",
    "ax.set_title('Distribution across time (boxplot per feature)')\n",
    "ax.grid(True, axis='y', linestyle='--', alpha=0.3)\n",
    "fig_box.tight_layout()\n",
    "fig_box.savefig('output/feature_boxplots.png', dpi=200)\n",
    "\n",
    "print(\"Saved: time_series_per_feature.png, heatmap_overview.png, feature_boxplots.png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b9bf68e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "import torchvision\n",
    "from lightning.pytorch import seed_everything\n",
    "import Loaders\n",
    "import numpy as np\n",
    "from pathlib import Path \n",
    "\n",
    "def setup_seed(seed):\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    np.random.seed(seed)\n",
    "    seed_everything(seed, workers=True)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "FOLD = 1\n",
    "setup_seed(2023)\n",
    "device = torch.device(f\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "df = pd.read_csv(\"./Dataset_RARP_video/dataset_videos_folds.csv\")\n",
    "\n",
    "ckpt_paths = [\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_0/checkpoints/RARP-epoch=20.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_1/checkpoints/RARP-epoch=32.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_2/checkpoints/RARP-epoch=28.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_3/checkpoints/RARP-epoch=27.ckpt\"),\n",
    "    Path(\"./log_XAblation_van_DINO/lightning_logs/version_4/checkpoints/RARP-epoch=30.ckpt\"),\n",
    "]\n",
    "\n",
    "train_set = df.loc[df[f\"Fold_{FOLD}\"] == \"train\"].sort_values(by=[\"label\", \"case\"]).to_dict(orient=\"records\")\n",
    "\n",
    "traintransformT2 = torch.nn.Sequential(\n",
    "    transforms.CenterCrop(300),\n",
    "    transforms.Resize((224, 224), antialias=True, interpolation=torchvision.transforms.InterpolationMode.BICUBIC),\n",
    "    transforms.RandomAffine(degrees=(-15, 15), scale=(0.8, 1.2), fill=0),\n",
    "    transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1),\n",
    "    transforms.GaussianBlur(kernel_size=3),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    ").to(device)\n",
    "\n",
    "traintransform_frame = torch.nn.Sequential(\n",
    "    transforms.RandomApply([\n",
    "        transforms.Lambda(lambda x: x + 0.01 * torch.randn_like(x)),\n",
    "        transforms.RandomErasing(1.0, value=0)\n",
    "        ], 0.3) #small noise\n",
    ").to(device)\n",
    "\n",
    "testVal_transform = torch.nn.Sequential(\n",
    "    transforms.CenterCrop(300),\n",
    "    transforms.Resize((224, 224), antialias=True, interpolation=torchvision.transforms.InterpolationMode.BICUBIC),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    ").to(device)\n",
    "\n",
    "key_frame_transform = torch.nn.Sequential(\n",
    "        transforms.Resize(256, antialias=True, interpolation=transforms.InterpolationMode.BICUBIC),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.Normalize([30.38144216, 42.03988769, 97.8896116], [40.63141752, 44.26910074, 50.29294373])\n",
    "    ).to(device)\n",
    "\n",
    "train_dataset = Loaders.RARP_Windowed_Video_frames_Dataset(\n",
    "    train_set, \n",
    "    train_mode=True, \n",
    "    window_length=64, \n",
    "    transform=traintransformT2, \n",
    "    transform_frame=traintransform_frame, \n",
    "    k_windows=1,\n",
    "    key_frames=True,\n",
    "    key_frame_transform=key_frame_transform\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df520225",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ff07c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "key_frame = train_dataset[3][4]\n",
    "key_frame = key_frame.repeat(1, 1, 1, 1)\n",
    "key_frame = key_frame.to(device, dtype=torch.float)\n",
    "key_frame.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ddbe594",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from Models import RARP_NVB_DINO_MultiTask\n",
    "Hybrid_TS = RARP_NVB_DINO_MultiTask.load_from_checkpoint(ckpt_paths[FOLD])\n",
    "Hybrid_TS.eval()\n",
    "\n",
    "with torch.no_grad():\n",
    "    pred, _, _ = Hybrid_TS(key_frame)\n",
    "    img_features = torch.cat((Hybrid_TS.last_conv_output_S, Hybrid_TS.last_conv_output_T), dim=1)\n",
    "    img_features = torch.nn.functional.adaptive_avg_pool2d(img_features, (1,1)).flatten(1) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58031f82",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred[0].item(), img_features[0].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a7b62a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "frame_idx = [0, 7, 15, 23, 31, 63] \n",
    "index = 3 \n",
    "frames = train_dataset[index][0][frame_idx]\n",
    "key_frame = train_dataset[index][4]\n",
    "\n",
    "Mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)\n",
    "Std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)\n",
    "\n",
    "frames = frames * Std + Mean\n",
    "frames = torch.clamp(frames, 0, 1)\n",
    "print(frames.shape)\n",
    "\n",
    "Mean = torch.tensor([30.38144216, 42.03988769, 97.8896116]).view(3,1,1)\n",
    "Std = torch.tensor([40.63141752, 44.26910074, 50.29294373]).view(3,1,1)\n",
    "key_frame = key_frame * Std + Mean\n",
    "key_frame = torch.clamp(key_frame[[2, 1, 0]] / 255, 0, 1)\n",
    "print(key_frame.shape)\n",
    "\n",
    "grid = torchvision.utils.make_grid(frames, nrow=len(frame_idx))\n",
    "\n",
    "fig, ax = plt.subplots(2, 1, figsize=(15, 8))\n",
    "\n",
    "ax[0].imshow(grid.permute(1,2,0).cpu().numpy())\n",
    "ax[0].axis(\"off\")\n",
    "\n",
    "ax[1].imshow(key_frame.permute(1,2,0).cpu().numpy())\n",
    "ax[1].axis(\"off\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ef02eab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "\n",
    "img = cv2.imread(\"Dataset_RARP_video/188/frames/frame_0003.webp\", cv2.IMREAD_COLOR)\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "img = torch.from_numpy(img.transpose((2, 0, 1)))\n",
    "img = img.float() / 255\n",
    "\n",
    "img = torchvision.transforms.functional.center_crop(img, 300)\n",
    "img = torchvision.transforms.functional.resize(img, (224, 224), antialias=True, interpolation=torchvision.transforms.InterpolationMode.BICUBIC)\n",
    "\n",
    "img =  torch.clamp(img, 0, 1)\n",
    "\n",
    "\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.imshow(img.permute(1,2,0).cpu().numpy())\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9d37eac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "\n",
    "dataset_video = Path(\"./Dataset_RARP_video/dataset_videos_folds.csv\")\n",
    "dataset_key_frame = Path(\"./dataset.csv\")\n",
    "\n",
    "db = pd.read_csv(dataset_video)\n",
    "dbk = pd.read_csv(dataset_key_frame)\n",
    "\n",
    "dbk = dbk[[\"case\", \"image_path\"]].rename(columns={\"image_path\": \"key_frame\"})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22ebde02",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = db.merge(dbk, on=\"case\", how=\"left\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73cdb3c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "db.to_csv(dataset_video, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "319fb566",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# load csv\n",
    "df = pd.read_csv(\"./test_reports/FOLD_0/mil_test_predictions_epoch38.csv\")\n",
    "\n",
    "# parameters (must match dataset)\n",
    "W = 10\n",
    "stride = 125          # seconds\n",
    "window_len = 72       # seconds\n",
    "\n",
    "window_starts = np.arange(W) * stride          # [0, 125, 250, ...]\n",
    "window_centers_sec = window_starts + window_len / 2\n",
    "window_centers_min = window_centers_sec / 60.0\n",
    "\n",
    "# choose one case\n",
    "row = df.iloc[0]\n",
    "\n",
    "alpha = np.array([row[f\"alpha_{i:02d}\"] for i in range(W)])\n",
    "\n",
    "plt.figure(figsize=(8, 3))\n",
    "plt.plot(window_centers_min, alpha, marker=\"o\")\n",
    "plt.xlabel(\"Time in surgery (minutes)\")\n",
    "plt.ylabel(\"MIL attention\")\n",
    "plt.title(\n",
    "    f\"Case {row['case_id']} | \"\n",
    "    f\"P(NVB)={row['prob']:.2f} | GT={row['y_true']}\"\n",
    ")\n",
    "plt.grid(True)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d57dd97a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 1.5))\n",
    "plt.imshow(alpha[np.newaxis, :], aspect=\"auto\", cmap=\"jet\")\n",
    "plt.yticks([])\n",
    "plt.xticks(\n",
    "    ticks=np.arange(W),\n",
    "    labels=[f\"{t:.1f}\" for t in window_centers_min],\n",
    ")\n",
    "plt.xlabel(\"Time in surgery (minutes, window centers)\")\n",
    "plt.colorbar(label=\"MIL attention\")\n",
    "plt.title(f\"Temporal Attention Map | Case {row['case_id']}\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cac66b06",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 2))\n",
    "\n",
    "for i in range(W):\n",
    "    start = window_starts[i] / 60.0\n",
    "    end   = (window_starts[i] + window_len) / 60.0\n",
    "    plt.hlines(\n",
    "        y=alpha[i],\n",
    "        xmin=start,\n",
    "        xmax=end,\n",
    "        linewidth=6,\n",
    "    )\n",
    "\n",
    "plt.xlabel(\"Time in surgery (minutes)\")\n",
    "plt.ylabel(\"MIL attention\")\n",
    "plt.title(f\"Window-level attention coverage | Case {row['case_id']}\")\n",
    "plt.grid(True)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e8f69306",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "40 % 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3286680e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 990224 characters]",
      "text/plain": [
       "<Figure size 2000x4500 with 44 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# load csv\n",
    "df = pd.read_csv(\"./test_reports_XT5/FOLD_1/mil_test_predictions_epoch47.csv\")\n",
    "th = 0.5\n",
    "\n",
    "# parameters (must match dataset)\n",
    "W = 10\n",
    "stride = 125          # seconds\n",
    "window_len = 72       # seconds\n",
    "\n",
    "window_starts = np.arange(W) * stride          # [0, 125, 250, ...]\n",
    "window_centers_sec = window_starts + window_len / 2\n",
    "window_centers_min = window_centers_sec / 60.0\n",
    "\n",
    "lenRow = len(df) \n",
    "lenRow += 0 if (lenRow % 2) == 0 else 1\n",
    "row_count = lenRow // 2\n",
    "\n",
    "fig, ax = plt.subplots(row_count, 2, figsize=(20, 45))\n",
    "# choose one case\n",
    "c = -1\n",
    "for r, row in df.iterrows():\n",
    "    plot_r = r % row_count\n",
    "    c += 1 if ( r % row_count ) == 0 else 0\n",
    "    \n",
    "    alpha = np.array([row[f\"alpha_{i:02d}\"] for i in range(W)])\n",
    "    \n",
    "    pred = \"NO_NVB\" if row['prob'] < 0.5 else \"NVB\"\n",
    "    label_txt = \"NVB\" if row[\"y_true\"] else \"NO_NVB\"\n",
    "    \n",
    "    for i in range(W):\n",
    "        start = window_starts[i] / 60.0\n",
    "        end   = (window_starts[i] + window_len) / 60.0\n",
    "        ax[plot_r, c].hlines(\n",
    "            y=alpha[i],\n",
    "            xmin=start,\n",
    "            xmax=end,\n",
    "            linewidth=6, color = (\"blue\" if pred == label_txt else \"red\" )\n",
    "        )\n",
    "    \n",
    "    ax[plot_r, c].set_xlabel(\"Time in surgery (minutes)\")\n",
    "    ax[plot_r, c].set_ylabel(\"MIL attention\")\n",
    "    ax[plot_r, c].set_title(f\"Window-level attention coverage | Case {row['case_id']} | P(x)={row['prob']:.4f} | Pred: {pred} | Label: {label_txt}\")\n",
    "    ax[plot_r, c].grid(True)\n",
    "    \n",
    "    \n",
    "    \n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d4b15e2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def plot_cluster_timeline(\n",
    "    csv_path: str,\n",
    "    case_id,\n",
    "    L_sec: int = 72,              # window length in seconds (e.g., 72 frames @ 1fps)\n",
    "    T_min: float = 20.0,          # total clip length in minutes\n",
    "    out_path: str | None = None,\n",
    "):\n",
    "    \"\"\"\n",
    "    CSV columns expected:\n",
    "    [case_id, win_idx, win_start_sec, label_nvb, win_start_min, cluster]\n",
    "    \"\"\"\n",
    "\n",
    "    df = pd.read_csv(csv_path)\n",
    "\n",
    "    # filter case\n",
    "    d = df[df[\"case_id\"].astype(str) == str(case_id)].copy()\n",
    "    if d.empty:\n",
    "        raise ValueError(f\"No rows found for case_id={case_id}\")\n",
    "\n",
    "    # sort windows by time\n",
    "    d = d.sort_values([\"win_start_min\", \"win_idx\"]).reset_index(drop=True)\n",
    "\n",
    "    # duration (minutes)\n",
    "    dur_min = L_sec / 60.0\n",
    "    d[\"dur_min\"] = dur_min\n",
    "\n",
    "    # stable cluster ordering\n",
    "    clusters = d[\"cluster\"].unique()\n",
    "    try:\n",
    "        clusters_sorted = sorted(clusters, key=lambda x: float(x))\n",
    "    except Exception:\n",
    "        clusters_sorted = sorted(clusters, key=lambda x: str(x))\n",
    "\n",
    "    cluster_to_y = {c: i for i, c in enumerate(clusters_sorted)}\n",
    "\n",
    "    # plot\n",
    "    fig, ax = plt.subplots(figsize=(12, 3.2))\n",
    "    bar_h = 0.8\n",
    "\n",
    "    for _, r in d.iterrows():\n",
    "        y = cluster_to_y[r[\"cluster\"]]\n",
    "        ax.broken_barh([(r[\"win_start_min\"], r[\"dur_min\"])], (y - bar_h / 2, bar_h), color=\"green\")\n",
    "\n",
    "    ax.set_title(f\"Cluster timeline | Case {case_id}\")\n",
    "    ax.set_xlabel(\"Time in surgery (minutes)\")\n",
    "    ax.set_ylabel(\"Cluster\")\n",
    "    ax.set_yticks(list(cluster_to_y.values()))\n",
    "    ax.set_yticklabels([str(c) for c in clusters_sorted])\n",
    "\n",
    "    ax.set_xlim(0, T_min)\n",
    "    ax.grid(True, axis=\"x\", alpha=0.3)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    if out_path:\n",
    "        plt.savefig(out_path, dpi=200, bbox_inches=\"tight\")\n",
    "        print(f\"Saved: {out_path}\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8f658051",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19608 characters]",
      "text/plain": [
       "<Figure size 1200x320 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_cluster_timeline(\"hybrid_t-s_cluster_out_F1K3_40/windows_clusters.csv\", case_id=531)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2bc0d89c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cluster_timeline_single_row(csv_path: str, case_id, L_sec: int = 72, T_min: float = 20.0):\n",
    "    df = pd.read_csv(csv_path)\n",
    "    d = df[df[\"case_id\"].astype(str) == str(case_id)].copy()\n",
    "    if d.empty:\n",
    "        raise ValueError(f\"No rows found for case_id={case_id}\")\n",
    "    d = d.sort_values([\"win_start_min\", \"win_idx\"]).reset_index(drop=True)\n",
    "\n",
    "    dur_min = L_sec / 60.0\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(12, 1.8))\n",
    "    y0, h = 0.0, 0.8\n",
    "\n",
    "    for _, r in d.iterrows():\n",
    "        ax.broken_barh([(r[\"win_start_min\"], dur_min)], (y0, h))\n",
    "        ax.text(r[\"win_start_min\"] + dur_min/2, y0 + h + 0.05, str(r[\"cluster\"]),\n",
    "                ha=\"center\", va=\"bottom\", fontsize=8)\n",
    "\n",
    "    ax.set_title(f\"Cluster timeline (single row) | Case {case_id}\")\n",
    "    ax.set_xlabel(\"Time in surgery (minutes)\")\n",
    "    ax.set_yticks([])\n",
    "    ax.set_xlim(0, T_min)\n",
    "    ax.grid(True, axis=\"x\", alpha=0.3)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "68c48754",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18588 characters]",
      "text/plain": [
       "<Figure size 1200x180 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_cluster_timeline_single_row(\"cluster_out_F1K3/windows_clusters.csv\", case_id=531)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9da01b5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
