{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "13:30:10 [I] klustakwik KlustaKwik2 version 0.2.6\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import expipe\n",
    "import pathlib\n",
    "import numpy as np\n",
    "import spatial_maps.stats as stats\n",
    "import septum_mec\n",
    "import septum_mec.analysis.data_processing as dp\n",
    "import septum_mec.analysis.registration\n",
    "import head_direction.head as head\n",
    "import spatial_maps as sp\n",
    "import speed_cells.speed as spd\n",
    "import re\n",
    "import joblib\n",
    "import multiprocessing\n",
    "import shutil\n",
    "import psutil\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "# import seaborn as sns\n",
    "from distutils.dir_util import copy_tree\n",
    "from neo import SpikeTrain\n",
    "import scipy\n",
    "\n",
    "from tqdm.notebook import tqdm_notebook as tqdm\n",
    "tqdm.pandas()\n",
    "\n",
    "from spike_statistics.core import permutation_resampling\n",
    "\n",
    "from spikewaveform.core import calculate_waveform_features_from_template, cluster_waveform_features\n",
    "\n",
    "from septum_mec.analysis.plotting import violinplot, despine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.rc('axes', titlesize=12)\n",
    "plt.rcParams.update({\n",
    "    'font.size': 12, \n",
    "    'figure.figsize': (6, 4), \n",
    "    'figure.dpi': 150\n",
    "})\n",
    "\n",
    "output_path = pathlib.Path(\"output\") / \"cell-count\"\n",
    "(output_path / \"statistics\").mkdir(exist_ok=True, parents=True)\n",
    "(output_path / \"figures\").mkdir(exist_ok=True, parents=True)\n",
    "output_path.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_loader = dp.Data()\n",
    "actions = data_loader.actions\n",
    "project = data_loader.project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "identification_action = actions['identify-neurons']\n",
    "sessions = pd.read_csv(identification_action.data_path('sessions'))\n",
    "units = pd.read_csv(identification_action.data_path('units'))\n",
    "session_units = pd.merge(sessions, units, on='action')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "stim_action = actions['stimulus-response']\n",
    "stim_results = pd.read_csv(stim_action.data_path('results'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "statistics_action = actions['calculate-statistics']\n",
    "shuffling = actions['shuffling']\n",
    "\n",
    "statistics_results = pd.read_csv(statistics_action.data_path('results'))\n",
    "statistics_results = session_units.merge(statistics_results, how='left')\n",
    "quantiles_95 = pd.read_csv(shuffling.data_path('quantiles_95'))\n",
    "action_columns = ['action', 'channel_group', 'unit_name']\n",
    "data = pd.merge(statistics_results, quantiles_95, on=action_columns, suffixes=(\"\", \"_threshold\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['unit_day'] = data.apply(lambda x: str(x.unit_idnum) + '_' + x.action.split('-')[1], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.merge(stim_results, how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform_action = actions['waveform-analysis']\n",
    "waveform_results = pd.read_csv(waveform_action.data_path('results')).drop('template', axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data.merge(waveform_results, how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = ['#d95f02','#e7298a']\n",
    "labels = ['11 Hz', '30 HZ']\n",
    "queries = ['frequency==11', 'frequency==30']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.bs = data.bs.astype(bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.loc[data.eval('t_i_peak == t_i_peak and not bs'), 'ns_inhibited'] = True\n",
    "data.ns_inhibited.fillna(False, inplace=True)\n",
    "\n",
    "data.loc[data.eval('t_i_peak != t_i_peak and not bs'), 'ns_not_inhibited'] = True\n",
    "data.ns_not_inhibited.fillna(False, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of sessions above threshold 194\n",
      "Number of animals 4\n"
     ]
    }
   ],
   "source": [
    "query = (\n",
    "    'gridness > gridness_threshold and '\n",
    "    'information_rate > information_rate_threshold and '\n",
    "    'gridness > .2 and '\n",
    "    'average_rate < 25'\n",
    ")\n",
    "sessions_above_threshold = data.query(query)\n",
    "print(\"Number of sessions above threshold\", len(sessions_above_threshold))\n",
    "print(\"Number of animals\", len(sessions_above_threshold.groupby(['entity'])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of gridcells 139\n",
      "Number of gridcell recordings 231\n",
      "Number of animals 4\n"
     ]
    }
   ],
   "source": [
    "gridcell_sessions = data[data.unit_day.isin(sessions_above_threshold.unit_day.values)]\n",
    "print(\"Number of gridcells\", gridcell_sessions.unit_idnum.nunique())\n",
    "print(\"Number of gridcell recordings\", len(gridcell_sessions))\n",
    "print(\"Number of animals\", len(gridcell_sessions.groupby(['entity'])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.loc[:,'gridcell'] = np.nan\n",
    "data['gridcell'] = data.isin(gridcell_sessions)\n",
    "\n",
    "data.loc[data.eval('not gridcell and bs'), 'bs_not_gridcell'] = True\n",
    "data.bs_not_gridcell.fillna(False, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Gridcell</th>\n",
       "      <th>BS not gridcell</th>\n",
       "      <th>NSi</th>\n",
       "      <th>NSni</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>entity</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1833</th>\n",
       "      <td>94</td>\n",
       "      <td>165</td>\n",
       "      <td>16</td>\n",
       "      <td>52</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1834</th>\n",
       "      <td>14</td>\n",
       "      <td>216</td>\n",
       "      <td>4</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1839</th>\n",
       "      <td>19</td>\n",
       "      <td>70</td>\n",
       "      <td>11</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1849</th>\n",
       "      <td>8</td>\n",
       "      <td>229</td>\n",
       "      <td>8</td>\n",
       "      <td>23</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        Gridcell  BS not gridcell  NSi  NSni\n",
       "entity                                      \n",
       "1833          94              165   16    52\n",
       "1834          14              216    4     7\n",
       "1839          19               70   11     5\n",
       "1849           8              229    8    23"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "table = pd.DataFrame()\n",
    "table['Gridcell'] = data.drop_duplicates('unit_idnum').query('gridcell').groupby('entity')['action'].count()\n",
    "table['BS not gridcell'] = data.drop_duplicates('unit_idnum').query('bs_not_gridcell').groupby('entity')['action'].count()\n",
    "table['NSi'] = data.drop_duplicates('unit_idnum').query('ns_inhibited').groupby('entity')['action'].count()\n",
    "table['NSni'] = data.drop_duplicates('unit_idnum').query('ns_not_inhibited').groupby('entity')['action'].count()\n",
    "table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Store results in Expipe action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "action = project.require_action(\"cell-count\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/media/storage/expipe/septum-mec/actions/stimulus-response/data/data/times.feather',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/data/psth.feather',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-bs_not_gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_not_inhibited.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-bs_not_gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-nsi-ns.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-nsi-ns-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_not_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_not_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_inhibited.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_not_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_not_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_not_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-bs_not_gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-bs_not_gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-bs_not_gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-bs_not_gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_not_inhibited.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-nsi-ns-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_inhibited.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-bs_not_gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-bs_not_gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-gc-ns.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-bs_not_gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_not_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_not_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-nsi-ns.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-gc-ns-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-gc-ns.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_not_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_not_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-bs_not_gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-bs_not_gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-bs_not_gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-bs_not_gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/response-probability-gc-ns-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-bs_not_gridcell.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_inhibited.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_not_inhibited.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_not_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-bs_not_gridcell-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_not_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_inhibited-stim-mec.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-gridcell.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_e_peak-ns_not_inhibited.svg',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-bs_not_gridcell-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_e_peak-ns_not_inhibited.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-p_i_peak-ns_inhibited-stim-mec.png',\n",
       " '/media/storage/expipe/septum-mec/actions/stimulus-response/data/figures/histogram-t_i_peak-ns_inhibited.png']"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "copy_tree(output_path, str(action.data_path()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "septum_mec.analysis.registration.store_notebook(action, \"20_cell_count.ipynb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}