{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "0ca78190",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import nibabel as nib\n",
    "import pandas as pd\n",
    "\n",
    "output_dir = Path(\"/media/storage/yalab-dev/voxelops/\")\n",
    "output_dir.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "id": "715fb349",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from templateflow import api as tf\n",
    "\n",
    "NAMING_SCHEME = \"Schaefer2018N{n_parcels}n{n_networks}Tian2020S{scale}\"\n",
    "\n",
    "\n",
    "def load_lut(n_schaefer: int | str, scale: int | str):\n",
    "    n_4s = int(n_schaefer) + 56\n",
    "    atlas_4s = output_dir / \"atlases\" / f\"atlas-4S{n_4s}Parcels\"\n",
    "    parcels_4s = atlas_4s / f\"atlas-4S{n_4s}Parcels_dseg.tsv\"\n",
    "    parcels_4s = pd.read_csv(parcels_4s, sep=\"\\t\")\n",
    "    parcels_4s[\"atlas_name\"] = parcels_4s[\"atlas_name\"].replace(\n",
    "        {f\"4S{n_4s}\": f\"Schaefer2018N{n_schaefer}n7\"}\n",
    "    )\n",
    "    return parcels_4s.iloc[:n_schaefer, :]\n",
    "\n",
    "\n",
    "def combine_atlases(schaefer_path, tian_path):\n",
    "    schaefer_img = nib.load(schaefer_path)\n",
    "    schaefer_data = np.round(schaefer_img.get_fdata()).astype(int)\n",
    "    tian_img = nib.load(tian_path)\n",
    "    tian_data = np.round(tian_img.get_fdata()).astype(int)\n",
    "    combined_data = schaefer_data.copy()\n",
    "    max_label = int(np.max(schaefer_data))\n",
    "    tian_data_nonzero = tian_data > 0\n",
    "    combined_data[tian_data_nonzero] = tian_data[tian_data_nonzero] + max_label\n",
    "    # generate image with integer data type\n",
    "    combined_img = nib.Nifti1Image(\n",
    "        combined_data, schaefer_img.affine, schaefer_img.header, dtype=np.int16\n",
    "    )\n",
    "    return combined_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60c5b033",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing atlas: Schaefer2018N100n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N100n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N100n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N100n7Tian2020S4\n",
      "Processing atlas: Schaefer2018N200n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N200n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N200n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N200n7Tian2020S4\n",
      "Processing atlas: Schaefer2018N300n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N300n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N300n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N300n7Tian2020S4\n",
      "Processing atlas: Schaefer2018N400n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N400n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N400n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N400n7Tian2020S4\n",
      "Processing atlas: Schaefer2018N500n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N500n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N500n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N500n7Tian2020S4\n",
      "Processing atlas: Schaefer2018N600n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N600n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N600n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N600n7Tian2020S4\n",
      "Processing atlas: Schaefer2018N700n7Tian2020S1\n",
      "Schaefer atlas with 700 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N700n7Tian2020S2\n",
      "Schaefer atlas with 700 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N700n7Tian2020S3\n",
      "Schaefer atlas with 700 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N700n7Tian2020S4\n",
      "Schaefer atlas with 700 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N800n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N800n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N800n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N800n7Tian2020S4\n",
      "Processing atlas: Schaefer2018N900n7Tian2020S1\n",
      "Schaefer atlas with 900 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N900n7Tian2020S2\n",
      "Schaefer atlas with 900 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N900n7Tian2020S3\n",
      "Schaefer atlas with 900 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N900n7Tian2020S4\n",
      "Schaefer atlas with 900 parcels not found in TemplateFlow. Skipping.\n",
      "Processing atlas: Schaefer2018N1000n7Tian2020S1\n",
      "Processing atlas: Schaefer2018N1000n7Tian2020S2\n",
      "Processing atlas: Schaefer2018N1000n7Tian2020S3\n",
      "Processing atlas: Schaefer2018N1000n7Tian2020S4\n"
     ]
    }
   ],
   "source": [
    "tian_schaefer_dir = Path(\n",
    "    \"/media/storage/yalab-dev/Tian2020MSA_v1.4/Tian2020MSA/3T/Subcortex-Only\"\n",
    ")\n",
    "atlas_list = [\n",
    "    f.name\n",
    "    for f in tian_schaefer_dir.glob(\"MNIvolumetric/*MNI152NLin2009cAsym_1mm.nii.gz\")\n",
    "]\n",
    "\n",
    "target_lut_format = (\n",
    "    \"Schaefer2018_{n_parcels}Parcels_7Networks_order_Tian_Subcortex_{scale}_label.txt\"\n",
    ")\n",
    "\n",
    "tian_fillers = {\n",
    "    \"network_label\": \"Subcortex\",\n",
    "}\n",
    "\n",
    "\n",
    "def extract_atlas_info(atlas_name: str):\n",
    "    parts = atlas_name.split(\"_\")\n",
    "    n_schaefer = parts[1].replace(\"Parcels\", \"\")\n",
    "    scale = parts[-4]\n",
    "    return n_schaefer, scale\n",
    "\n",
    "\n",
    "new_atlases = []\n",
    "\n",
    "for n_schaefer in np.arange(100, 1001, 100):\n",
    "    for scale in [1, 2, 3, 4]:\n",
    "        output_atlas_name = NAMING_SCHEME.format(\n",
    "            n_parcels=n_schaefer, n_networks=7, scale=scale\n",
    "        )\n",
    "        new_atlases.append(output_atlas_name)\n",
    "        print(f\"Processing atlas: {output_atlas_name}\")\n",
    "        output_atlas_dir = (\n",
    "            output_dir / \"Schaefer2018Tian2020_atlases\" / f\"atlas-{output_atlas_name}\"\n",
    "        )\n",
    "        output_atlas_dir.mkdir(parents=True, exist_ok=True)\n",
    "        output_tsv = output_atlas_dir / f\"atlas-{output_atlas_name}_dseg.tsv\"\n",
    "        output_nii = (\n",
    "            output_atlas_dir\n",
    "            / f\"atlas-{output_atlas_name}_space-MNI152NLin2009cAsym_res-01_dseg.nii.gz\"\n",
    "        )\n",
    "        # if output_tsv.exists() and output_nii.exists():\n",
    "        #     print(f\"Atlas {output_atlas_name} already exists. Skipping.\")\n",
    "        #     continue\n",
    "        schaefer_nifti = tf.get(\n",
    "            \"MNI152NLin2009cAsym\",\n",
    "            atlas=\"Schaefer2018\",\n",
    "            desc=f\"{n_schaefer}Parcels7Networks\",\n",
    "            resolution=\"01\",\n",
    "            extension=\".nii.gz\",\n",
    "        )\n",
    "        if not schaefer_nifti:\n",
    "            print(\n",
    "                f\"Schaefer atlas with {n_schaefer} parcels not found in TemplateFlow. Skipping.\"\n",
    "            )\n",
    "            continue\n",
    "        parcels_4s = load_lut(n_schaefer, scale)\n",
    "        tian_nii_path = (\n",
    "            tian_schaefer_dir / f\"Tian_Subcortex_S{scale}_3T_2009cAsym_1mm.nii.gz\"\n",
    "        )\n",
    "        tian_txt_path = tian_schaefer_dir / f\"Tian_Subcortex_S{scale}_3T_label.txt\"\n",
    "        df_tian = pd.read_csv(tian_txt_path, header=None).rename(columns={0: \"label\"})\n",
    "        # concat the two dataframes on the label column\n",
    "        merged_df = pd.concat([parcels_4s, df_tian], axis=0).reset_index(drop=True)\n",
    "        # Fill in missing values with the tian_fillers\n",
    "        for key, value in tian_fillers.items():\n",
    "            merged_df[key] = merged_df[key].fillna(value)\n",
    "        merged_df[\"atlas_name\"] = merged_df[\"atlas_name\"].fillna(f\"Tian2020S{scale}\")\n",
    "        merged_df[\"index\"] = merged_df.index + 1\n",
    "        # now create the combined atlas nifti\n",
    "\n",
    "        combined_img = combine_atlases(schaefer_nifti, tian_nii_path)\n",
    "        # save df to tsv\n",
    "        merged_df.to_csv(output_tsv, sep=\"\\t\", index=False)\n",
    "\n",
    "        nib.save(combined_img, output_nii)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "d5438144",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "d = Path(\"/media/storage/yalab-dev/voxelops/Schaefer2018Tian2020_atlases\")\n",
    "for atlas_dir in d.iterdir():\n",
    "    if atlas_dir.is_dir() and atlas_dir.name.startswith(\"Schaefer\"):\n",
    "        new_name = f\"atlas-{atlas_dir.name}\"\n",
    "        # rename directory\n",
    "        atlas_dir.rename(d / new_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "VoxelOps",
   "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
