{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Obtain Cortical Layers in Channels\n",
    "\n",
    "This file is used to obtain the cortical layers information for a channels dataframe. The original dataframe only contains information on the visual areas, but since we have the common coordinates framework (CCF), we can get the cortical lauyers by using the Allen CCFv3 mouse atlas. This code is an adaptation from the code used in the work from Siegle et al (Siegle, J.H., Jia, X., Durand, S. et al. Survey of spiking in the mouse visual system reveals functional hierarchy. Nature 592, 86–92 (2021). https://doi.org/10.1038/s41586-020-03171-x)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import os\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import nrrd\n",
    "import re\n",
    "from allensdk.brain_observatory.ecephys.ecephys_project_cache import EcephysProjectCache\n",
    "#import files needed\n",
    "dir = os.getcwd()\n",
    "manifest_path = os.path.join(dir, \"manifest.json\")\n",
    "cache = EcephysProjectCache.from_warehouse(manifest=manifest_path)\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "streamlines, header = nrrd.read(dir+'\\ccf_volumes\\laplacian_10.nrrd')\n",
    "annotations = np.load(dir+'/ccf_volumes/annotation_volume_10um_by_index.npy')\n",
    "structure_tree = pd.read_csv(dir+'/ccf_volumes/ccf_structure_tree.csv', index_col=0)\n",
    "\n",
    "def get_layer_name(acronym):\n",
    "    #this \n",
    "    try:\n",
    "        layer = int(re.findall(r'\\d+', acronym)[0])\n",
    "        if layer == 3:\n",
    "            layer = 0\n",
    "        return layer\n",
    "    except IndexError:\n",
    "        return 0\n",
    "    \n",
    "def get_structure_ids(df, annotations):\n",
    "    \n",
    "    x = (df.anterior_posterior_ccf_coordinate.values / 10).astype('int')\n",
    "    y = (df.dorsal_ventral_ccf_coordinate.values / 10).astype('int')\n",
    "    z = (df.left_right_ccf_coordinate.values / 10).astype('int')\n",
    "    \n",
    "    x[x < 0] = 0\n",
    "    y[y < 0] = 0\n",
    "    z[z < 0] = 0\n",
    "    \n",
    "    structure_ids = annotations[x, y, z] - 1 # annotation volume is Matlab-indexed\n",
    "    \n",
    "    return structure_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total channels: 23\n",
      "1    3\n",
      "2    6\n",
      "4    3\n",
      "5    5\n",
      "6    6\n",
      "Name: cortical_layer, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "ses = 16 #session number\n",
    "\n",
    "#load the channels\n",
    "with open('channels/chans_V1_'+str(ses)+'.pickle', 'rb') as f:\n",
    "    chans_V1 = pickle.load(f)\n",
    "\n",
    "#if chans_V1 already has the column 'cortical_layer' we raise an error\n",
    "if 'cortical_layer' in chans_V1.columns:\n",
    "    raise ValueError('cortical_layer already in chans_V1')\n",
    "\n",
    "#we obtain the CCf coordinates for each channel and adapt them to the CCFv3 notation\n",
    "x = (chans_V1.anterior_posterior_ccf_coordinate.values / 10).astype('int')\n",
    "y = (chans_V1.dorsal_ventral_ccf_coordinate.values / 10).astype('int')\n",
    "z = (chans_V1.left_right_ccf_coordinate.values / 10).astype('int')\n",
    "cortical_depth = streamlines[x, y, z]\n",
    "chans_V1['cortical_depth'] = 0\n",
    "chans_V1.loc[chans_V1.anterior_posterior_ccf_coordinate > 0, 'cortical_depth'] = cortical_depth\n",
    "\n",
    "structure_ids = get_structure_ids(chans_V1, annotations)\n",
    "structure_acronyms = structure_tree.loc[structure_ids].acronym\n",
    "layers = [get_layer_name(acronym) for acronym in structure_acronyms]     \n",
    "chans_V1['cortical_layer'] = layers\n",
    "\n",
    "#count the number of channels for each cortical_layer instance\n",
    "# Assuming 'df' is your DataFrame\n",
    "layer_counts = chans_V1['cortical_layer'].value_counts()\n",
    "layer_counts = layer_counts.sort_index()\n",
    "print('total channels:', len(chans_V1))\n",
    "print(layer_counts) #where the left column is the layer and the right column is the count\n",
    "\n",
    "#save chans_V1 as pickle\n",
    "with open('channels/chans_V1_'+str(ses)+'.pickle', 'wb') as f:\n",
    "    pickle.dump(chans_V1, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "allensdk",
   "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}