{
 "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": [
      "08:31:41 [I] klustakwik KlustaKwik2 version 0.2.6\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import spatial_maps as sp\n",
    "import septum_mec.analysis.data_processing as dp\n",
    "import septum_mec.analysis.registration\n",
    "import expipe\n",
    "import os\n",
    "import pathlib\n",
    "import numpy as np\n",
    "import exdir\n",
    "import pandas as pd\n",
    "import optogenetics as og\n",
    "import quantities as pq\n",
    "import shutil\n",
    "from distutils.dir_util import copy_tree\n",
    "\n",
    "from septum_mec.analysis.stimulus_response import stimulus_response_latency, compute_response\n",
    "\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from tqdm._tqdm_notebook import tqdm_notebook\n",
    "tqdm_notebook.pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "std_gaussian_kde = 0.04\n",
    "window_size = 0.03"
   ]
  },
  {
   "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": [
    "output = pathlib.Path('output/stimulus-response')\n",
    "(output / 'figures').mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "identify_neurons = actions['identify-neurons']\n",
    "units = pd.read_csv(identify_neurons.data_path('units'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def process(row):\n",
    "    \n",
    "    action_id = row['action']\n",
    "    channel_id = int(row['channel_group'])\n",
    "    unit_id = int(row['unit_name'])    \n",
    "    \n",
    "    spike_times = data_loader.spike_train(action_id, channel_id, unit_id)\n",
    "    \n",
    "    spike_times = np.array(spike_times)\n",
    "    \n",
    "    stim_times = data_loader.stim_times(action_id)\n",
    "    \n",
    "    nan_series = pd.Series({\n",
    "            't_e_peak': np.nan,\n",
    "            'p_e_peak': np.nan,\n",
    "            't_i_peak': np.nan,\n",
    "            'p_i_peak': np.nan\n",
    "        })\n",
    "    \n",
    "    if stim_times is None:\n",
    "        return nan_series\n",
    "    \n",
    "    stim_times = np.array(stim_times)\n",
    "    \n",
    "    times, spikes, kernel, p_e, p_i = stimulus_response_latency(\n",
    "        spike_times, stim_times, window_size, std_gaussian_kde)\n",
    "    \n",
    "    # if no spikes detected after stimulus nan is returned\n",
    "    if all(np.isnan([p_e, p_i])):\n",
    "        return nan_series\n",
    "        \n",
    "    t_e_peak, p_e_peak, t_i_peak, p_i_peak = compute_response(\n",
    "        spike_times, stim_times, times, kernel, p_e, p_i)\n",
    "\n",
    "    return pd.Series({\n",
    "        't_e_peak': t_e_peak,\n",
    "        'p_e_peak': p_e_peak,\n",
    "        't_i_peak': t_i_peak,\n",
    "        'p_i_peak': p_i_peak\n",
    "    })\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3f0400c75c745c789907bd763bf2833",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1281), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mikkel/apps/expipe-project/septum-mec/septum_mec/analysis/stimulus_response.py:33: RuntimeWarning: invalid value encountered in less\n",
      "  if any(times[idxs_i] < te_peak):\n"
     ]
    }
   ],
   "source": [
    "results = units.merge(\n",
    "    units.progress_apply(process, axis=1), \n",
    "    left_index=True, right_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "results.loc[:, ['t_e_peak', 't_i_peak', 'p_e_peak', 'p_i_peak']].hist()\n",
    "plt.gcf().savefig(output / 'figures' / 'summary_histogram.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save to expipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action = project.require_action(\"stimulus-response\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action.modules['parameters'] = {\n",
    "    'window_size': window_size,\n",
    "    'std_gaussian_kde': std_gaussian_kde\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action.data['results'] = 'results.csv'\n",
    "results.to_csv(action.data_path('results'), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stuff = {\n",
    "    \"figures\": \"figures\",\n",
    "#     \"statistics\": \"statistics\"\n",
    "}\n",
    "\n",
    "for key, value in stuff.items():\n",
    "    action.data[key] = value\n",
    "    data_path = action.data_path(key)\n",
    "    data_path.parent.mkdir(exist_ok=True, parents=True)\n",
    "    source = output / value\n",
    "    if source.is_file():\n",
    "        shutil.copy(source, data_path)\n",
    "    else:\n",
    "        copy_tree(str(source), str(data_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "septum_mec.analysis.registration.store_notebook(action, \"10-calculate-stimulus-response.ipynb\")"
   ]
  }
 ],
 "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": 2
}