{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8149a94",
   "metadata": {},
   "source": [
    "# General Linear Model (GLM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b54b132a",
   "metadata": {},
   "source": [
    "In this tutorial, which is adapted from the Nilearn docs, we will go through a simple workflow of the first level general linear modeling with a BIDS dataset from openneuro. This analysis is only performed on **one** subject.\n",
    "\n",
    "This tutorial is based on the [Nilearn GLM tutorial](https://nilearn.github.io/stable/auto_examples/04_glm_first_level/plot_bids_features.html#sphx-glr-auto-examples-04-glm-first-level-plot-bids-features-py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f514ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8313a041",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "Import packages that will be used globally and set up output directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72d1dfdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "import sys \n",
    "if not sys.warnoptions:\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    \n",
    "import os\n",
    "import typing as ty\n",
    "from pathlib import Path\n",
    "\n",
    "from pydra.compose import python, workflow\n",
    "from pydra.engine.submitter import Submitter\n",
    "from fileformats.generic import File, Directory\n",
    "from fileformats.text import Csv\n",
    "import pandas as pd\n",
    "from scipy.stats import norm\n",
    "\n",
    "import nibabel as nib\n",
    "# These functions were removed within nilearn, so this notebook needs to be rewritten\n",
    "# to use the 'openneuro' module instead\n",
    "# from nilearn.datasets import (\n",
    "#         fetch_openneuro_dataset_index,\n",
    "#         fetch_openneuro_dataset,\n",
    "#         select_from_index,\n",
    "#     )\n",
    "from nilearn.interfaces.fsl import get_design_from_fslmat\n",
    "from nilearn.glm.first_level import first_level_from_bids\n",
    "from nilearn.reporting import get_clusters_table, make_glm_report\n",
    "from nilearn.plotting import (\n",
    "    plot_glass_brain,\n",
    "    plot_img_comparison,\n",
    "    plot_contrast_matrix,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5716cb50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get current directory\n",
    "pydra_tutorial_dir = os.path.dirname(os.getcwd())\n",
    "\n",
    "# set up output directory\n",
    "workflow_dir = Path(pydra_tutorial_dir) / 'outputs'\n",
    "workflow_out_dir = workflow_dir / '6_glm'\n",
    "\n",
    "# create the output directory if not exit\n",
    "os.makedirs(workflow_out_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1878928b",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_out_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cafd6a1",
   "metadata": {},
   "source": [
    "## Create tasks\n",
    "\n",
    "In this section, we converte major steps into tasks.\n",
    "Each pydra task can have multiple python functions. We recommend to put those logically more related functions into the same task.\n",
    "\n",
    "It is very **important** to keep in mind what adjacent tasks of your current task will be.\n",
    "1. Your previous task will decide your arguments in the current task\n",
    "2. Your next task will be impacted by the returns in the current task"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "823780ab",
   "metadata": {},
   "source": [
    "### fetch openneuro BIDS dataset\n",
    "\n",
    "In this task, we do the following:\n",
    "1. get openneuro dataset index\n",
    "2. specify exclusion patterns and number of subjects\n",
    "3. download the data we need\n",
    "\n",
    "\n",
    "**Notes:** Here we still use `n_subjects` as an argument. Given that we will only analyze one subject, you can also remove this argument and specify `n_subjects =1` in `select_from_index`. If you do, do not forget to modify the argument in the workflow later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2ab134c",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(outputs=[\"data_dir\"])\n",
    "def GetOpenneuroDataset(exclusion_patterns: list, n_subjects: int) -> str:\n",
    "    _, urls = fetch_openneuro_dataset_index()\n",
    "    urls = select_from_index(\n",
    "        urls, exclusion_filters=exclusion_patterns, n_subjects=n_subjects\n",
    "    )\n",
    "    data_dir, _ = fetch_openneuro_dataset(urls=urls)\n",
    "    return data_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b4899de",
   "metadata": {},
   "source": [
    "### obtain FirstLevelModel objects automatically and fit arguments\n",
    "\n",
    "To get the first level model(s) we have to specify\n",
    "1. the dataset directory\n",
    "2. the task_label\n",
    "3. the space_label\n",
    "4. the folder with the desired derivatives (fMRIPrep)\n",
    "\n",
    "In our case, we only have one subject so we will only have one first level model.\n",
    "Then, for this model, we will obtain\n",
    "1. the list of run images\n",
    "2. events\n",
    "3. confound regressors\n",
    "\n",
    "Those are inferred from the confounds.tsv files available in the BIDS dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c2710dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(outputs=[\"model\", \"imgs\", \"subject\"])\n",
    "def GetInfoFromBids(\n",
    "    data_dir: Directory,\n",
    "    task_label: str,\n",
    "    space_label: str,\n",
    "    smoothing_fwhm: float,\n",
    "    derivatives_folder: Directory,\n",
    ") -> ty.Tuple[ty.Any, list, str]:\n",
    "    (\n",
    "        models,\n",
    "        models_run_imgs,\n",
    "        models_events,\n",
    "        models_confounds,\n",
    "    ) = first_level_from_bids(\n",
    "        dataset_path=data_dir,\n",
    "        task_label=task_label,\n",
    "        space_label=space_label,\n",
    "        smoothing_fwhm=smoothing_fwhm,\n",
    "        derivatives_folder=derivatives_folder,\n",
    "    )\n",
    "    model, imgs, events, confounds = (\n",
    "        models[0],\n",
    "        models_run_imgs[0],\n",
    "        models_events[0],\n",
    "        models_confounds[0],\n",
    "    )\n",
    "    subject = 'sub-' + model.subject_label\n",
    "    return model, imgs, subject"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5af99cb",
   "metadata": {},
   "source": [
    "### Get design matrix\n",
    "\n",
    "This task does the following:\n",
    "1. read the design matrix in `.mat`\n",
    "2. rename the column\n",
    "3. save the new design matrix as `.csv`\n",
    "\n",
    "**Think:** What if we don't save the new design matrix, but `return` it directly? In other words, we `return` a `pandas.DataFrame` instead of a `path`. What will happen? Worth a try :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bdfcfd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(outputs=[\"dm_path\"])\n",
    "def GetDesignMatrix(data_dir: Directory, subject: str) -> Csv:\n",
    "    fsl_design_matrix_path = data_dir.joinpath(\n",
    "        'derivatives',\n",
    "        'task',\n",
    "        subject,\n",
    "        'stopsignal.feat',\n",
    "        'design.mat',\n",
    "    )\n",
    "    design_matrix = get_design_from_fslmat(\n",
    "        fsl_design_matrix_path, column_names=None\n",
    "    )\n",
    "\n",
    "    design_columns = [\n",
    "        'cond_%02d' % i for i in range(len(design_matrix.columns))\n",
    "    ]\n",
    "    design_columns[0] = 'Go'\n",
    "    design_columns[4] = 'StopSuccess'\n",
    "    design_matrix.columns = design_columns\n",
    "    dm_path = Path('designmatrix.csv')\n",
    "    design_matrix.to_csv(dm_path, index=None)\n",
    "    return dm_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1cb37d0",
   "metadata": {},
   "source": [
    "### Fit the first level model\n",
    "\n",
    "What we are doing here is:\n",
    "1. use the design matrix to fit the first level model\n",
    "2. compute the contrast\n",
    "3. save the z_map and masker for further use\n",
    "4. generate a glm report (HTML file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65cec504",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(outputs=[\"model\", \"z_map_path\", \"masker\", \"glm_report_file\"])\n",
    "def ModelFit(model, imgs, dm_path, contrast: str) -> ty.Tuple[ty.Any, str, ty.Any, str]:\n",
    "    design_matrix = pd.read_csv(dm_path)\n",
    "    model.fit(imgs, design_matrices=[design_matrix])\n",
    "    z_map = model.compute_contrast(contrast)\n",
    "    z_map_path = Path('firstlevel_z_map.nii.gz')\n",
    "    z_map.to_filename(z_map_path)\n",
    "    masker_path = Path('firstlevel_masker.nii.gz')\n",
    "    masker = model.masker_\n",
    "    glm_report_file = Path('glm_report.html')\n",
    "    report = make_glm_report(model, contrast)\n",
    "    report.save_as_html(glm_report_file)\n",
    "    return model, z_map_path, masker, glm_report_file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05576ba4",
   "metadata": {},
   "source": [
    "### Get cluster table\n",
    "\n",
    "For publication purposes, we obtain a cluster table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4a86a6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(outputs=[\"output_file\"])\n",
    "def ClusterTable(z_map_path: File) -> Csv:\n",
    "    stat_img = nib.load(z_map_path)\n",
    "    output_file = Path('cluster_table.csv')\n",
    "    df = get_clusters_table(\n",
    "        stat_img, stat_threshold=norm.isf(0.001), cluster_threshold=10\n",
    "    )\n",
    "    df.to_csv(output_file, index=None)\n",
    "    return output_file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e8effd",
   "metadata": {},
   "source": [
    "### Make plots\n",
    "\n",
    "Here we want to make some plots to display our results and compare the result from FSL.\n",
    "1. plot nilearn z-map\n",
    "2. plot fsl z-map\n",
    "3. plot nilearn and fsl comparison\n",
    "4. plot design matrix contrast\n",
    "\n",
    "You can also separate this task into multiple sub-tasks. But it makes more sense to put them into one task as they use the same files and function `nilearn.plotting` repeatedly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0f78107",
   "metadata": {},
   "outputs": [],
   "source": [
    "@python.define(outputs=[\"output_file1\", \"output_file2\", \"output_file3\", \"output_file4\"])\n",
    "def Plots(\n",
    "    data_dir: Directory,\n",
    "    dm_path: File,\n",
    "    z_map_path: File,\n",
    "    contrast: str,\n",
    "    subject: str,\n",
    "    masker\n",
    ") -> ty.Tuple[str, str, str, str]:\n",
    "    # plot and save nilearn z-map\n",
    "    z_map = nib.load(z_map_path)\n",
    "    output_file1 = Path('nilearn_z_map.jpg')\n",
    "    plot_glass_brain(\n",
    "        z_map,\n",
    "        output_file=output_file1,\n",
    "        colorbar=True,\n",
    "        threshold=norm.isf(0.001),\n",
    "        title='Nilearn Z map of \"StopSuccess - Go\" (unc p<0.001)',\n",
    "        plot_abs=False,\n",
    "        display_mode='ortho',\n",
    "    )\n",
    "\n",
    "    # plot and save fsl z-map\n",
    "    fsl_z_map = nib.load(\n",
    "        os.path.join(\n",
    "            data_dir,\n",
    "            'derivatives',\n",
    "            'task',\n",
    "            subject,\n",
    "            'stopsignal.feat',\n",
    "            'stats',\n",
    "            'zstat12.nii.gz',\n",
    "        )\n",
    "    )\n",
    "    output_file2 = Path('fsl_z_map.jpg')\n",
    "    plot_glass_brain(\n",
    "        fsl_z_map,\n",
    "        output_file=output_file2,\n",
    "        colorbar=True,\n",
    "        threshold=norm.isf(0.001),\n",
    "        title='FSL Z map of \"StopSuccess - Go\" (unc p<0.001)',\n",
    "        plot_abs=False,\n",
    "        display_mode='ortho',\n",
    "    )\n",
    "\n",
    "    # plot and save nilearn and fsl comparison\n",
    "    plot_img_comparison(\n",
    "        [z_map],\n",
    "        [fsl_z_map],\n",
    "        masker,\n",
    "        output_dir=workflow_out_dir,\n",
    "        ref_label='Nilearn',\n",
    "        src_label='FSL',\n",
    "    )\n",
    "    old = Path('0000.png')\n",
    "    new = Path('nilearn_fsl_comp.jpg')\n",
    "    os.rename(old, new)\n",
    "    output_file3 = new\n",
    "    print(output_file3)\n",
    "\n",
    "    # plot and save design matrix contrast\n",
    "    design_matrix = pd.read_csv(dm_path)\n",
    "    output_file4 = Path('firstlevel_contrast.jpg')\n",
    "    plot_contrast_matrix(contrast, design_matrix, output_file=output_file4)\n",
    "    return output_file1, output_file2, output_file3, output_file4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12a99b96",
   "metadata": {},
   "source": [
    "## Make a workflow from tasks\n",
    "\n",
    "Now we have created all tasks we need for this first level analysis, and there are two choices for our next step.\n",
    "1. create one workflow to connect all tasks together\n",
    "2. create sub-workflows with some closely related tasks, and connect these workflows along with other tasks into a larger workflow.\n",
    "\n",
    "We recommend the second approach as it is always a good practice to group tasks, especially when there are a large number of tasks in the analysis.\n",
    "\n",
    "Our analysis can be divided into three parts: (1) get/read the data, (2) analyze the data, and (3) plot the result, where (1) and (3) only have one task each. So we can put all tasks in (2) into one workflow and name it as `firstlevel` or whatever you prefer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e79e9b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@workflow.define(outputs=[\"z_map\", \"masker\", \"subject\", \"dm_path\", \"cluster_table\", \"glm_report\"])\n",
    "def FirstLevelWorkflow(\n",
    "    data_dir: Directory,\n",
    "    contrast: str,\n",
    "    output_dir: Path,\n",
    "    task_label: str = 'stopsignal',\n",
    "    space_label: str = 'MNI152NLin2009cAsym',\n",
    "    derivatives_folder: str = 'derivatives/fmriprep',\n",
    "    smoothing_fwhm: float = 5.0,\n",
    ") -> ty.Tuple[str, str, str, File, str, str]:\n",
    "\n",
    "    # add task - get_info_from_bids\n",
    "    get_info_from_bids = workflow.add(\n",
    "        GetInfoFromBids(\n",
    "            data_dir=data_dir,\n",
    "            task_label=task_label,\n",
    "            space_label=space_label,\n",
    "            derivatives_folder=derivatives_folder,\n",
    "            smoothing_fwhm=smoothing_fwhm,\n",
    "        )\n",
    "    )\n",
    "    # add task - get_designmatrix\n",
    "    get_designmatrix = workflow.add(\n",
    "        GetDesignMatrix(\n",
    "            data_dir=data_dir,\n",
    "            subject=get_info_from_bids.subject,\n",
    "        )\n",
    "    )\n",
    "    l1estimation = workflow.add(\n",
    "        ModelFit(\n",
    "            model=get_info_from_bids.model,\n",
    "            imgs=get_info_from_bids.imgs,\n",
    "            dm_path=get_designmatrix.dm_path,\n",
    "            contrast=contrast,\n",
    "        )\n",
    "    )\n",
    "    # add task - cluster_table\n",
    "    cluster_table = workflow.add(\n",
    "        ClusterTable(\n",
    "            z_map_path=l1estimation.z_map_path,\n",
    "        )\n",
    "    )\n",
    "    # specify output\n",
    "    return (\n",
    "        l1estimation.z_map_path,\n",
    "        l1estimation.masker,\n",
    "        get_info_from_bids.subject,\n",
    "        get_designmatrix.dm_path,\n",
    "        cluster_table.output_file,\n",
    "        l1estimation.glm_report_file,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "657690ea",
   "metadata": {},
   "source": [
    "## The overaching workflow\n",
    "\n",
    "Connect other tasks and the above workflow into one\n",
    "\n",
    "Now we need to create the overaching glm workflow that connects the above workflow and other tasks (e.g., `get/read the data` and `plot the result`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d055c5d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@workflow.define(outputs=[\"output1\", \"output2\", \"output3\", \"output4\"])\n",
    "def FullWorkflow(\n",
    "    output_dir: Path,\n",
    "    n_subjects: int = 1,\n",
    "    contrast: str = 'StopSuccess - Go',\n",
    "    exclusion_patterns: list[str] | None = None,\n",
    ") -> tuple[ty.Any, ty.Any, ty.Any, ty.Any]:\n",
    "    if exclusion_patterns is None:\n",
    "        exclusion_patterns = [\n",
    "            '*group*',\n",
    "            '*phenotype*',\n",
    "            '*mriqc*',\n",
    "            '*parameter_plots*',\n",
    "            '*physio_plots*',\n",
    "            '*space-fsaverage*',\n",
    "            '*space-T1w*',\n",
    "            '*dwi*',\n",
    "            '*beh*',\n",
    "            '*task-bart*',\n",
    "            '*task-rest*',\n",
    "            '*task-scap*',\n",
    "            '*task-task*',\n",
    "        ]\n",
    "\n",
    "    get_openneuro_dataset = workflow.add(\n",
    "        GetOpenneuroDataset(\n",
    "            exclusion_patterns=exclusion_patterns,\n",
    "            n_subjects=n_subjects,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    wf_firstlevel = workflow.add(\n",
    "        FirstLevelWorkflow(\n",
    "            data_dir=get_openneuro_dataset.data_dir,\n",
    "            contrast=contrast,\n",
    "            output_dir=output_dir,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    plots = workflow.add(\n",
    "        Plots(\n",
    "            data_dir=get_openneuro_dataset.data_dir,\n",
    "            dm_path=wf_firstlevel.dm_path,\n",
    "            z_map_path=wf_firstlevel.z_map,\n",
    "            contrast=contrast,\n",
    "            subject=wf_firstlevel.subject,\n",
    "            masker=wf_firstlevel.masker,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return (\n",
    "        plots.output_file1,\n",
    "        plots.output_file2,\n",
    "        plots.output_file3,\n",
    "        plots.output_file4,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2e9a46",
   "metadata": {},
   "source": [
    "## Run Workflow Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a90088e",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "wf = FullWorkflow(output_dir=workflow_out_dir, n_subjects=1, contrast='StopSuccess - Go')\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    with Submitter(worker='cf', n_procs=4) as sub:\n",
    "        results = sub(wf)\n",
    "\n",
    "    print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f540cdd4",
   "metadata": {},
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8def869",
   "metadata": {},
   "source": [
    "If you arrive here without any errors, yay, you just made your first pydra workflow for a first-level GLM!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b0585e3",
   "metadata": {},
   "source": [
    "## Examine folder structure\n",
    "\n",
    "Let's take a look at what you have got."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75c1cfc9",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "! ls ../outputs/6_glm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56aeee0c",
   "metadata": {},
   "source": [
    "### Plot figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f657571",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "\n",
    "if not results.errored:\n",
    "    # First-level contrast\n",
    "    Image(filename='../outputs/6_glm/firstlevel_contrast.jpg')\n",
    "\n",
    "    # Nilearn Z map\n",
    "    Image(filename='../outputs/6_glm/nilearn_z_map.jpg')\n",
    "\n",
    "    # FSL Z map\n",
    "    Image(filename='../outputs/6_glm/fsl_z_map.jpg')\n",
    "\n",
    "    # Nilearn and FSL comparison\n",
    "    Image(filename='../outputs/6_glm/nilearn_fsl_comp.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "081bf13a",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3d55272",
   "metadata": {},
   "source": [
    "What if we need to run the first-level GLM on multiple subject? We will need the `splitter`.\n",
    "\n",
    "So, where should we add `.split`?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wf12",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}