{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# robreg\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "# Allow unsupported MPS ops to fall back to CPU so the notebook can still test MPS.\n",
    "os.environ.setdefault(\"PYTORCH_ENABLE_MPS_FALLBACK\", \"1\")\n",
    "\n",
    "import nibabel as nib\n",
    "import torch\n",
    "\n",
    "REPO_ROOT = Path.cwd().resolve()\n",
    "if not (REPO_ROOT / \"neuroreg\").exists():\n",
    "    REPO_ROOT = REPO_ROOT.parent\n",
    "if str(REPO_ROOT) not in sys.path:\n",
    "    sys.path.insert(0, str(REPO_ROOT))\n",
    "\n",
    "from neuroreg import robreg\n",
    "from neuroreg.image import map\n",
    "from neuroreg.image.map import map_r2r\n",
    "from neuroreg.transforms import LTA, affine_dist\n",
    "from neuroreg.transforms.matrices import get_affine\n",
    "\n",
    "DATA_DIR = Path(\".\")\n",
    "orig_path = DATA_DIR / \"140_orig.mgz\"\n",
    "trg_path = DATA_DIR / \"140_trg.mgz\"\n",
    "reg_path = DATA_DIR / \"140_reg.mgz\"\n",
    "lta_path = DATA_DIR / \"140.lta\"\n",
    "gt_lta_path = DATA_DIR / \"140_ground_truth.lta\"\n",
    "\n",
    "\n",
    "device = \"cpu\"\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = \"mps\"\n",
    "    print(\"MPS detected; unsupported ops will fall back to CPU via PYTORCH_ENABLE_MPS_FALLBACK=1.\")\n",
    "\n",
    "print(f\"Running on device: {device}\")\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Download a test image.\n",
    "url = \"https://surfer.nmr.mgh.harvard.edu/pub/data/tutorial_data/buckner_data/tutorial_subjs/140/mri/orig.mgz\"\n",
    "if not orig_path.exists():\n",
    "    urlretrieve(url, orig_path)\n",
    "else:\n",
    "    print(f\"File {orig_path} already exists, skipping download\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Load the image and create a small synthetic rigid transform.\n",
    "img = nib.load(str(orig_path))\n",
    "idata = torch.from_numpy(img.get_fdata()).float()\n",
    "img_affine = torch.from_numpy(img.affine).float()\n",
    "\n",
    "translation_mm = torch.tensor([3.1, 0.6, 1.2])\n",
    "rotation_rad = torch.tensor([0.08, 0.04, 0.02])\n",
    "v2v_gt = get_affine(translation=translation_mm, rotvec=rotation_rad).float()\n",
    "r2r_gt = img_affine @ v2v_gt @ torch.linalg.inv(img_affine)\n",
    "\n",
    "print(f\"Applied Vox-to-Vox matrix:\\n{v2v_gt}\")\n",
    "print(f\"Ground-truth RAS-to-RAS matrix:\\n{r2r_gt}\")\n",
    "\n",
    "# Map the image and save the moved volume.\n",
    "mapped_data = map(idata, transform=v2v_gt, is_torch_mat=False)\n",
    "target_img = nib.MGHImage(mapped_data.squeeze().numpy(), img.affine, img.header)\n",
    "target_img.to_filename(str(trg_path))\n",
    "\n",
    "# Save the ground-truth transform as an LTA.\n",
    "LTA.from_matrix(r2r_gt.numpy(), str(orig_path), img, str(trg_path), target_img).write(str(gt_lta_path))\n",
    "print(f\"Saved synthetic target image to: {trg_path}\")\n",
    "print(f\"Saved ground-truth LTA to: {gt_lta_path}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Register the synthetic pair.\n",
    "# By default, robreg returns an RAS-to-RAS transform.\n",
    "Mr2r = robreg(\n",
    "    str(orig_path),\n",
    "    str(trg_path),\n",
    "    device=device,\n",
    ")\n",
    "Mr2r_cpu = Mr2r.detach().cpu()\n",
    "\n",
    "recovered_data = map_r2r(\n",
    "    idata,\n",
    "    Mr2r_cpu.float(),\n",
    "    source_affine=img_affine,\n",
    "    target_affine=img_affine,\n",
    "    target_shape=tuple(int(v) for v in img.shape[:3]),\n",
    "    mode=\"bilinear\",\n",
    ")\n",
    "recovered_img = nib.MGHImage(recovered_data.detach().cpu().numpy(), img.affine, img.header)\n",
    "recovered_img.to_filename(str(reg_path))\n",
    "LTA.from_matrix(Mr2r_cpu.numpy(), str(orig_path), img, str(trg_path), target_img).write(str(lta_path))\n",
    "\n",
    "recovered_lta = LTA.read(str(lta_path))\n",
    "affine_distance = float(affine_dist(Mr2r_cpu.float(), r2r_gt.float(), radius=100.0))\n",
    "saved_lta_distance = float(\n",
    "    affine_dist(torch.from_numpy(recovered_lta.r2r()).float(), r2r_gt.float(), radius=100.0)\n",
    ")\n",
    "\n",
    "print(f\"\\nRecovered RAS-to-RAS matrix:\\n{Mr2r_cpu}\")\n",
    "print(f\"\\nSaved LTA RAS-to-RAS matrix:\\n{torch.from_numpy(recovered_lta.r2r()).float()}\")\n",
    "print(f\"\\nAffine distance to the ground-truth transform (r=100 mm): {affine_distance:.4f}\")\n",
    "print(f\"Affine distance for the saved LTA (r=100 mm): {saved_lta_distance:.4f}\")\n",
    "print(f\"Mapped output saved to: {reg_path}\")\n",
    "print(f\"Recovered LTA saved to: {lta_path}\")\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}