{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7dbf62e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "sys.path.append(os.path.join(os.getcwd(), '..'))\n",
    "\n",
    "import numpy as np\n",
    "import h5py\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from scipy import signal\n",
    "from scipy import stats\n",
    "from target import build_tgt_matrix\n",
    "import pandas as pd\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "220a45cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.OutputArea.prototype._should_scroll = function(lines) {\n",
       "    return false;\n",
       "}\n"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%javascript\n",
    "IPython.OutputArea.prototype._should_scroll = function(lines) {\n",
    "    return false;\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "06dca5bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "source = '/home/sobolev/nevermind/Andrey/data'\n",
    "report = '/home/sobolev/nevermind/Andrey/analysis/PSTH'\n",
    "\n",
    "selected_sessions = [\n",
    "'009266_hippoSIT_2023-04-17_17-04-17',  # ch17, 20 + 55 correction, 5067 events. Showcase for N2 / N3 mod in target\n",
    "'009266_hippoSIT_2023-04-18_10-10-37',  # ch17, 10 + 55 correction, 5682 events\n",
    "'009266_hippoSIT_2023-04-18_17-03-10',  # ch17, 6 + 55 correction, 5494 events: FIXME very weird 1-2nd in target, find out\n",
    "'009266_hippoSIT_2023-04-19_10-33-51',  # ch17, 4 + 55 correction, 6424 events: very weird 1-2nd in target, find out\n",
    "'009266_hippoSIT_2023-04-20_08-57-39',  # ch1, 1 + 55 correction, 6424 events. Showcase for N2 / N3 mod in target\n",
    "'009266_hippoSIT_2023-04-24_16-56-55',  # ch17, 5 + 55* correction, 6165 events, frequency\n",
    "'009266_hippoSIT_2023-04-26_08-20-17',  # ch17, 12 + 55* correction, 6095 events, duration - showcase for N2 \n",
    "'009266_hippoSIT_2023-05-02_12-22-14',  # ch20, 10 + 55 correction, 5976 events, FIXME very weird 1-2nd in target, find out\n",
    "'009266_hippoSIT_2023-05-04_09-11-06',  # ch17, 5 + 55* correction, 4487 events, coma session with baseline AEPs\n",
    "'009266_hippoSIT_2023-05-04_19-47-15',  # ch20, 2 + 55 correction, 5678 events, duration\n",
    "]\n",
    "\n",
    "session = selected_sessions[0]\n",
    "\n",
    "animal      = session.split('_')[0]\n",
    "sessionpath = os.path.join(source, animal, session)\n",
    "aeps_file   = os.path.join(sessionpath, 'AEPs.h5')\n",
    "h5name      = os.path.join(sessionpath, session + '.h5')\n",
    "report_path = os.path.join(report, session)\n",
    "if not os.path.exists(report_path):\n",
    "    os.makedirs(report_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b5a09f6a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((5067, 200), (73, 5))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with h5py.File(h5name, 'r') as f:\n",
    "    tl = np.array(f['processed']['timeline'])  # time, X, Y, speed, etc.\n",
    "    trials = np.array(f['processed']['trial_idxs'])  # t_start_idx, t_end_idx, x_tgt, y_tgt, r_tgt, result\n",
    "    cfg = json.loads(f['processed'].attrs['parameters'])\n",
    "    \n",
    "with h5py.File(aeps_file, 'r') as f:\n",
    "    aeps = np.array(f['aeps'])\n",
    "    aeps_events = np.array(f['aeps_events'])\n",
    "    \n",
    "# TODO find better way. Remove outliers\n",
    "aeps[aeps > 5000]  =  5000\n",
    "aeps[aeps < -5000] = -5000\n",
    "\n",
    "# load metrics\n",
    "AEP_metrics_lims = {}\n",
    "AEP_metrics_raw  = {}\n",
    "AEP_metrics_norm = {}\n",
    "with h5py.File(aeps_file, 'r') as f:\n",
    "    for metric_name in f['raw']:\n",
    "        AEP_metrics_raw[metric_name]  = np.array(f['raw'][metric_name])\n",
    "        AEP_metrics_norm[metric_name] = np.array(f['norm'][metric_name])\n",
    "        AEP_metrics_lims[metric_name] = [int(x) for x in f['raw'][metric_name].attrs['limits'].split(',')]\n",
    "\n",
    "tgt_dur = cfg['experiment']['target_duration']\n",
    "tgt_matrix = build_tgt_matrix(tl, aeps_events, cfg['experiment']['target_duration'])\n",
    "\n",
    "aeps.shape, tgt_matrix.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0e36f45",
   "metadata": {},
   "source": [
    "## Is performance dependent on AEP states?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "9294c349",
   "metadata": {},
   "outputs": [],
   "source": [
    "# separate states based on AEP metrics - metric mean before entering\n",
    "def compute_state_idxs(metric_name, n_pulses=10):\n",
    "    idxs_aeps_high, idxs_aeps_low = [], []  # indices to tgt_matrix\n",
    "\n",
    "    metric_mean = AEP_metrics_norm[metric_name].mean()\n",
    "    for i, tgt_enter in enumerate(tgt_matrix):\n",
    "        if tgt_enter[2] - n_pulses < 0:\n",
    "            continue\n",
    "        metric_inst = AEP_metrics_norm[metric_name][tgt_enter[2] - n_pulses:tgt_enter[2]].mean()\n",
    "        if metric_inst > metric_mean:\n",
    "            idxs_aeps_high.append(i)\n",
    "        else:\n",
    "            idxs_aeps_low.append(i)\n",
    "            \n",
    "    return idxs_aeps_low, idxs_aeps_high"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "02ffa773",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'N1': (0.6164383561643836, 0.3835616438356164),\n",
       " 'P1': (0.5342465753424658, 0.4657534246575342),\n",
       " 'P2': (0.6164383561643836, 0.3835616438356164),\n",
       " 'P3': (0.3835616438356164, 0.6164383561643836)}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test if the animal will stay in the island depending on high / low AEP metrics\n",
    "\n",
    "metric_name = 'P1'\n",
    "n_pulses = 8\n",
    "predictions = {}\n",
    "\n",
    "for metric_name in AEP_metrics_norm.keys():\n",
    "    idxs_aeps_low, idxs_aeps_high = compute_state_idxs(metric_name, n_pulses)\n",
    "\n",
    "    actual = tgt_matrix[:, 4]\n",
    "    predicted = np.zeros(len(tgt_matrix))\n",
    "    predicted[idxs_aeps_high] = 1\n",
    "    high_low = (predicted == actual).sum() / len(tgt_matrix)\n",
    "    \n",
    "    predicted = np.zeros(len(tgt_matrix))\n",
    "    predicted[idxs_aeps_low] = 1\n",
    "    low_high = (predicted == actual).sum() / len(tgt_matrix)\n",
    "    \n",
    "    predictions[metric_name] = (high_low, low_high)\n",
    "    \n",
    "predictions"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}