{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# include controllers to the path\n",
    "import sys, os\n",
    "sys.path.append(os.getcwd())\n",
    "sys.path.append(os.path.join(os.getcwd(), 'controllers'))\n",
    "\n",
    "import cv2\n",
    "import threading\n",
    "import math\n",
    "import time\n",
    "import random\n",
    "import json\n",
    "import datetime\n",
    "import os, shutil\n",
    "import numpy as np\n",
    "import multiprocess as mp\n",
    "\n",
    "# controllers\n",
    "import nbimporter\n",
    "from controllers.situtils import FPSTimes\n",
    "from controllers.camera import WebcamStream\n",
    "from controllers.video import VideoWriter\n",
    "from controllers.position import PositionTracker\n",
    "from controllers.sound import SoundController\n",
    "from controllers.serial import MCSArduino, FakeArduino, Feeder\n",
    "from postprocessing import pack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load experiment settings\n",
    "\n",
    "For every experimental cofiguration you can copy the original 'settings.json' file, build your own specific experimental preset, save it in this folder as e.g. 'settings_elena.json' and load it here instead of 'settings.json'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg_filename = os.path.join('profiles', 'dardo_aSIT_single.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "with open(cfg_filename) as json_file:\n",
    "    cfg = json.load(json_file)\n",
    "cfg['experiment']['experiment_date'] = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')\n",
    "\n",
    "# print loaded settings\n",
    "#print(json.dumps(cfg, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check if the sound interface is there\n",
    "import sounddevice as sd\n",
    "asio = [x for x in sd.query_devices() if x['name'].find('ASIO') > 0]\n",
    "#if len(asio) == 0:\n",
    "#    raise SystemExit('The sound interface is not found. Please restart the computer')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize session folder\n",
    "\n",
    "Run the upcoming cell, to create a session folder and to save the chosen experimetal parameters to a JSON-file (\"experiment_id_parameters.json\"). The session folder will be created here where this notebook is located."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This session's protocols will be saved to this folder\n",
    "cfg_exp = cfg['experiment']\n",
    "experiment_id = \"%s_%s_%s\" % (cfg_exp['subject'], cfg_exp['experiment_type'], cfg_exp['experiment_date'])\n",
    "save_to = os.path.join('sessions', experiment_id)\n",
    "             \n",
    "if not os.path.exists(save_to):\n",
    "    os.makedirs(save_to)\n",
    "\n",
    "# update paths (assuming this paths are relative to this notebook)\n",
    "cfg['video']['file_path'] = os.path.join(save_to, 'video.avi')\n",
    "cfg['position']['file_path'] = os.path.join(save_to, 'positions.csv')\n",
    "cfg['experiment']['file_path'] = os.path.join(save_to, 'events.csv')\n",
    "cfg['sound']['file_path'] = os.path.join(save_to, 'sounds.csv')\n",
    "cfg['position']['background_light'] = os.path.join('assets', cfg['position']['background_light'])\n",
    "cfg['position']['background_dark'] = os.path.join('assets', cfg['position']['background_dark'])\n",
    "    \n",
    "# Saves all parameters to a JSON file with the user-defined \"Experiment ID\" as filename\n",
    "with open(os.path.join(save_to, experiment_id + '.json'), 'w') as f:\n",
    "    json.dump(cfg, f, indent=4)\n",
    "    \n",
    "with open(cfg['experiment']['file_path'], 'w') as f:\n",
    "    # state: 0 - trial start, 1 - trial success, 2 - trial fail\n",
    "    f.write('time, target_x, target_y, target_r, trial, state\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Island:    \n",
    "    def __init__(self, x, y, radius, sound_id, is_distractor=False):\n",
    "        self.x = x  # in meters\n",
    "        self.y = y  # in meters\n",
    "        self.r = radius  # in meters\n",
    "        self.sound_id = sound_id\n",
    "        self.is_distractor = is_distractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_islands(floor_radius, island_radius, distractors=0):\n",
    "    # all in meters\n",
    "    rho = np.sqrt(np.random.rand()) * (floor_radius - island_radius)\n",
    "    phi = np.random.rand() * 2 * np.pi\n",
    "    \n",
    "    # TODO write multi islands\n",
    "    return [Island(rho * np.sin(phi), rho * np.cos(phi), island_radius, 2)]  # always a list of Islands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def timeout(t_start):\n",
    "    return time.time() - t_start > cfg_exp['session_duration'] if t_start is not None else False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_event(*args):  # log start / end of a trial\n",
    "    with open(cfg_exp['file_path'], 'a') as f:\n",
    "        f.write(\",\".join([str(x) for x in args]) + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def switch_light(pt, board):\n",
    "    pt.switch_background()\n",
    "    board.switch_light()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start the experiment\n",
    "\n",
    "This cell contains code for animal tracking. We hope that the comments provided in the code suffice to understand the individual steps and to adjust them to your own setup and needs, if necessary.\n",
    "\n",
    "- press 's' to start recording\n",
    "- press 's' again to stop recording\n",
    "- press 'q' to quit\n",
    "\n",
    "The experiment will stop automatically if the pre-defined session duration is reached."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Webcam stream 1024.0:768.0 at 20.00 FPS started\n",
      "Position tracker stopped\n",
      "Video writer stopped\n",
      "Camera released\n"
     ]
    }
   ],
   "source": [
    "# actual sound selector: 0 - silence, 1 - foraging, 2 - target, 3 - distractor\n",
    "sound = mp.Value('i', 1)\n",
    "\n",
    "# experiment status: 1 - idle, 2 - running (recording, logging), 0 - stopped\n",
    "status = mp.Value('i', 1)\n",
    "\n",
    "# init the sync with the acquisition system via Arduino\n",
    "if cfg['experiment']['MCSArduinoPort'] == 'fake':\n",
    "    board = FakeArduino()\n",
    "else:\n",
    "    board = MCSArduino(cfg['experiment']['MCSArduinoPort'])\n",
    "\n",
    "# init the feeder\n",
    "feeder = Feeder(cfg['experiment']['feeder_port'])\n",
    "    \n",
    "# start the camera stream\n",
    "vs = WebcamStream(cfg['camera'])\n",
    "vs.start()\n",
    "\n",
    "# init video recorder\n",
    "vw = VideoWriter(status, vs, cfg['video'])\n",
    "vw.start()\n",
    "\n",
    "# start position tracking\n",
    "pt = PositionTracker(status, vs, cfg['position'])\n",
    "pt.start()\n",
    "\n",
    "# playing sound in a separate process for performance\n",
    "sc = mp.Process(target=SoundController.run, args=(sound, status, cfg['sound']))\n",
    "sc.start()\n",
    "\n",
    "timers = []\n",
    "fps = FPSTimes()\n",
    "names = ['camera', 'video', 'position', 'main']\n",
    "trial = 0\n",
    "rewards = 0\n",
    "t_start = None\n",
    "target_since = None\n",
    "punishment_since = None\n",
    "trial_start = time.time()\n",
    "phase = 0  # 0 - idle, 1 - foraging, 2 - inter-trial interval\n",
    "cfg_exp = cfg['experiment']\n",
    "cfg_pos = cfg['position']\n",
    "COLORS = {\n",
    "    'red': (0,0,255), 'green': (127,255,0), 'blue': (255,127,0), 'yellow': (0,127,255), \\\n",
    "    'black': (0,0,0), 'white': (255,255,255)\n",
    "}\n",
    "islands = []\n",
    "contour = None\n",
    "\n",
    "try:\n",
    "    while trial <= cfg_exp['trial_number'] and not timeout(t_start):\n",
    "        frame = vs.read()\n",
    "        if frame is None:\n",
    "            time.sleep(0.1)\n",
    "            continue # wait for the stream\n",
    "            \n",
    "        c_time = time.time()\n",
    "        fps.count()\n",
    "        status_color = COLORS['green'] if status.value == 1 else COLORS['red']\n",
    "\n",
    "        # -------- prepare the video frame ---------------\n",
    "        \n",
    "        # mask space outside arena\n",
    "        frame = cv2.bitwise_and(src1=frame, src2=pt.mask)\n",
    "        #frame = cv2.subtract(frame, pt.background)\n",
    "        #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "        # draw position center and contours\n",
    "        if pt.x_in_px is not None:\n",
    "            cv2.circle(frame, (pt.x_in_px, pt.y_in_px), 7, status_color, -1)\n",
    "            cv2.drawContours(frame, [pt.contour], 0, status_color, 1, cv2.LINE_AA)  \n",
    "            \n",
    "            contour = pt.contour\n",
    "\n",
    "        #l_s_x = cfg_pos['arena_x'] + cfg_pos['arena_radius'] - 20\n",
    "        #l_e_x = cfg_pos['arena_x'] + cfg_pos['arena_radius'] + 20\n",
    "        #cv2.line(frame, (l_s_x, cfg_pos['arena_y']), (l_e_x, cfg_pos['arena_y']), status_color, 5)\n",
    "            \n",
    "        # add FPS indicators\n",
    "        for i, ctrl in enumerate([vs, vw, pt, fps]):\n",
    "            cv2.putText(frame, '%s: %.2f FPS' % (names[i], ctrl.get_avg_fps()), \n",
    "                     (10, 30 + 20*(i+1)), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "\n",
    "        # size of the arena and status indicator\n",
    "        cv2.circle(frame, (cfg_pos['arena_x'], cfg_pos['arena_y']), cfg_pos['arena_radius'], COLORS['red'], 2)\n",
    "        cv2.circle(frame, (cfg_pos['arena_x'], cfg_pos['arena_y']), cfg_pos['floor_radius'], COLORS['red'], 2)\n",
    "        cv2.circle(frame, (20,20), 10, status_color, -6)\n",
    "\n",
    "        # draw islands\n",
    "        if len(islands) > 0:\n",
    "            for island in islands:\n",
    "                clr = COLORS['red'] if island.is_distractor else COLORS['green']\n",
    "                x, y = pt.meters_to_px(island.x, island.y)\n",
    "                cv2.circle(frame, (x, y), int(island.r/pt.pixel_size), clr, 2)\n",
    "\n",
    "        # positions of animal and target\n",
    "        if pt.x_in_px is not None:\n",
    "            cv2.putText(frame, 'Animal: %.3f %.3f' % (pt.x_in_m, pt.y_in_m), (10, 550), \\\n",
    "                        cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "            cv2.putText(frame, '%.2f %.2f' % (pt.x_in_m, pt.y_in_m), (pt.x_in_px + 10, pt.y_in_px + 10), \\\n",
    "                        cv2.FONT_HERSHEY_DUPLEX, .3, COLORS['white'])\n",
    "            \n",
    "        if len(islands) > 0:\n",
    "            target = [i for i in islands if not i.is_distractor][0]\n",
    "            cv2.putText(frame, 'Target: %.3f %.3f' % (target.x, target.y), (10, 570), \\\n",
    "                        cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "            \n",
    "        # stopwatch\n",
    "        stopwatch = 'Time: %.2f' % float(c_time - t_start) if t_start is not None else 'Time: Idle'\n",
    "        cv2.putText(frame, stopwatch, (10, 590), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "\n",
    "        # trial countdown (TODO add island countdown)\n",
    "        if status.value > 1:\n",
    "            text = 'Trial: %.2f' % float(cfg_exp['trial_duration'] - (c_time - trial_start))\n",
    "            cv2.putText(frame, text, (10, 610), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "            cv2.putText(frame, 'Trial: %s' % trial, (10, 630), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "            \n",
    "        # rewards\n",
    "        cv2.putText(frame, 'Rewards: %s' % rewards, (10, 650), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "        \n",
    "        # time since in target\n",
    "        if target_since is not None:\n",
    "            cv2.putText(frame, 'In target: %.2f' % float(c_time - target_since), (10, 670), cv2.FONT_HERSHEY_DUPLEX, .5, COLORS['white'])\n",
    "        \n",
    "        # Light or dark period\n",
    "        cv2.putText(frame, 'Light' if pt.is_light else 'Dark', (10, 690), cv2.FONT_HERSHEY_DUPLEX, .5, (255, 255, 255))            \n",
    "        \n",
    "        # assign the frame back to the video stream for other controllers\n",
    "        vs.frame_with_infos = frame\n",
    "        \n",
    "        cv2.imshow('Press (s)-to start/stop, (q)-to end', frame)\n",
    "\n",
    "        # -------- experiment logic ---------------\n",
    "        \n",
    "        # animals is either foraging (phase == 1) or in the inter-trial-interval (phase == 2)\n",
    "        if phase == 1: # foraging\n",
    "            if pt.x_in_m is not None and len(islands) > 0:\n",
    "                # check if animal in the island and for how long\n",
    "                tgt = [i for i in islands if not i.is_distractor][0]\n",
    "                distractors = [i for i in islands if i.is_distractor]\n",
    "                \n",
    "                if (pt.x_in_m - tgt.x)**2 + (pt.y_in_m - tgt.y)**2 <= cfg_exp['target_radius']**2:\n",
    "                    if target_since is None:  # just entered target\n",
    "                        target_since = c_time\n",
    "                        sound.value = 2\n",
    "\n",
    "                    elif c_time - target_since > cfg_exp['target_duration']: # successful trial\n",
    "                        log_event(c_time, round(tgt.x, 4), round(tgt.y, 4), round(tgt.r, 4), trial, 1) # log trial success\n",
    "                        feeder.feed()\n",
    "                        \n",
    "                        # init inter-trial interval and new trial\n",
    "                        trial_start = c_time + 10 + 5 * np.random.rand()  # random b/w 10-15 sec\n",
    "                        sound.value = 0  # silence\n",
    "                        islands = []\n",
    "                        phase = 2\n",
    "                        rewards += 1\n",
    "                        target_since = None\n",
    "                        \n",
    "                elif c_time - trial_start > cfg_exp['trial_duration']:  # trial failed and animal is not in the target\n",
    "                    log_event(c_time, round(tgt.x, 4), round(tgt.y, 4), round(tgt.r, 4), trial, 2) # log trial failed\n",
    "\n",
    "                    trial_start = c_time + 10 + 5 * np.random.rand()  # random b/w 10-15 sec\n",
    "                    sound.value = -1  # punishment\n",
    "                    punishment_since = c_time\n",
    "                    islands = []\n",
    "                    phase = 2\n",
    "                \n",
    "                else:\n",
    "                    target_since = None\n",
    "                    in_distractor = False\n",
    "                    \n",
    "                    for isl in distractors:  # maybe animal is in one of the distractors\n",
    "                        if (pt.x_in_m - isl.x)**2 + (pt.y_in_m - isl.y)**2 <= cfg_exp['target_radius']**2:\n",
    "                            sound.value = isl.sound_id\n",
    "                            in_distractor = True\n",
    "                    \n",
    "                    if not in_distractor:  # outside of the islands\n",
    "                        sound.value = 1\n",
    "            \n",
    "        elif phase == 2:  # inter-trial-interval\n",
    "            if punishment_since is not None and c_time - punishment_since > 10:  # max 10 sec of punishment\n",
    "                sound.value = 0\n",
    "                punishment_since = None\n",
    "            \n",
    "            if c_time > trial_start:\n",
    "                # init_new_trial\n",
    "                islands = generate_islands(cfg_pos['floor_r_in_meters'], cfg_exp['target_radius'])\n",
    "                sound.value = 1\n",
    "                phase = 1\n",
    "                trial += 1\n",
    "                \n",
    "                # log trial start\n",
    "                tgt = [i for i in islands if not i.is_distractor][0]\n",
    "                log_event(c_time, round(tgt.x, 4), round(tgt.y, 4), round(tgt.r, 4), trial, 0)\n",
    "        \n",
    "        # -------- key press events ---------------\n",
    "        \n",
    "        k = cv2.waitKey(33)\n",
    "        if k == ord('q'):\n",
    "            break\n",
    "\n",
    "        if k == ord('s'):\n",
    "            board.start_or_stop()  # start/stop data acquisition\n",
    "            c_time = time.time()   # more accurate time\n",
    "\n",
    "            if status.value == 1: # start the session\n",
    "                if t_start is None:\n",
    "                    t_start = c_time\n",
    "                trial_start = c_time\n",
    "                status.value = 2\n",
    "                \n",
    "                # init_new_trial\n",
    "                islands = generate_islands(cfg_pos['floor_r_in_meters'], cfg_exp['target_radius'])\n",
    "                sound.value = 1\n",
    "                phase = 1\n",
    "                trial += 1\n",
    "\n",
    "                # log trial start\n",
    "                tgt = [i for i in islands if not i.is_distractor][0]\n",
    "                log_event(c_time, round(tgt.x, 4), round(tgt.y, 4), round(tgt.r, 4), trial, 0)\n",
    "                \n",
    "                # init light events\n",
    "                timers = []\n",
    "                for event_t in cfg_exp['light_events']:\n",
    "                    timers.append(threading.Timer(event_t, switch_light, args=(pt, board)))                \n",
    "                for t in timers:\n",
    "                    t.start()\n",
    "                \n",
    "            elif status.value == 2:  # pause the session\n",
    "                if len(islands) > 0:\n",
    "                    tgt = [i for i in islands if not i.is_distractor][0]\n",
    "                    x, y, r = round(tgt.x, 4), round(tgt.y, 4), round(tgt.r, 4)\n",
    "                else:\n",
    "                    x, y, r = None, None, None\n",
    "                log_event(c_time, x, y, r, trial, -1)  # log experiment break\n",
    "                \n",
    "                status.value = 1\n",
    "                phase = 0\n",
    "                islands = []\n",
    "                for t in timers:\n",
    "                    t.cancel()\n",
    "                    \n",
    "        if k == ord('a'):\n",
    "            sound.value = 0 if sound.value == 1 else 1\n",
    "\n",
    "finally:\n",
    "    if status.value == 2:  # stop data acquisition, in case exited via timeout/trial number/quit\n",
    "        board.start_or_stop()\n",
    "        if len(islands) > 0:\n",
    "            tgt = [i for i in islands if not i.is_distractor][0]\n",
    "            x, y, r = round(tgt.x, 4), round(tgt.y, 4), round(tgt.r, 4)\n",
    "        else:\n",
    "            x, y, r = None, None, None\n",
    "        log_event(c_time, x, y, r, trial, -1)  # log experiment end\n",
    "        \n",
    "    status.value = 0\n",
    "    time.sleep(0.01)\n",
    "    \n",
    "    feeder.exit()\n",
    "    board.exit()\n",
    "    cv2.destroyAllWindows()\n",
    "    sc.join()\n",
    "    for ctrl in [pt, vw, vs]:\n",
    "        ctrl.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((-0.02564459930313589, -0.19874564459930313), 264)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, y = contour[0][0][0], contour[0][0][1]\n",
    "pt.px_to_meters(x, y), len(contour)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctr_in_m = np.array([pt.px_to_meters(x, y) for x, y in zip(contour[:, 0, 0], contour[:, 0, 1])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctr_in_px = np.array([x for x in zip(contour[:, 0, 0], contour[:, 0, 1])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "440"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "contour[:, 0, 0].min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x598a0d0>]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16300 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from scipy import stats\n",
    "\n",
    "x = ctr_in_px[:, 0]\n",
    "y = ctr_in_px[:, 1]\n",
    "\n",
    "slope, intercept, r_value, p_value, std_err = stats.linregress(x,y)\n",
    "x_line = np.linspace(x.min(), x.max(), 10)\n",
    "y_line = x_line * slope + intercept\n",
    "\n",
    "plt.scatter(x, y)\n",
    "plt.plot(x_line, y_line)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merge data in HDF5 file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "session_path = save_to\n",
    "#session_path = os.path.join('sessions', '2021-07-30_09-24-14')  # some particular session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "could not convert string to float: 'None'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-36-03bc734495ad>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;31m# do pack data to HDF5\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 5\u001b[1;33m \u001b[0mh5name\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mpack\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0msession_path\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32mD:\\runSIT\\postprocessing.ipynb\u001b[0m in \u001b[0;36mpack\u001b[1;34m(session_path)\u001b[0m\n\u001b[0;32m     33\u001b[0m     \u001b[1;34m\"    \\n\"\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     34\u001b[0m     \u001b[1;34m\"    File has the following structure:\\n\"\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 35\u001b[1;33m     \u001b[1;34m\"    \\n\"\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     36\u001b[0m     \u001b[1;34m\"    /raw\\n\"\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     37\u001b[0m     \u001b[1;34m\"        /positions      - raw positions from .csv\\n\"\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\.conda\\envs\\runsit\\lib\\site-packages\\numpy\\lib\\npyio.py\u001b[0m in \u001b[0;36mloadtxt\u001b[1;34m(fname, dtype, comments, delimiter, converters, skiprows, usecols, unpack, ndmin, encoding, max_rows, like)\u001b[0m\n\u001b[0;32m   1144\u001b[0m         \u001b[1;31m# converting the data\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1145\u001b[0m         \u001b[0mX\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1146\u001b[1;33m         \u001b[1;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mread_data\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0m_loadtxt_chunksize\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1147\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mX\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1148\u001b[0m                 \u001b[0mX\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0marray\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\.conda\\envs\\runsit\\lib\\site-packages\\numpy\\lib\\npyio.py\u001b[0m in \u001b[0;36mread_data\u001b[1;34m(chunk_size)\u001b[0m\n\u001b[0;32m    995\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    996\u001b[0m             \u001b[1;31m# Convert each value according to its column and store\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 997\u001b[1;33m             \u001b[0mitems\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mconv\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mval\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mconv\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mval\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mconverters\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mvals\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    998\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    999\u001b[0m             \u001b[1;31m# Then pack it according to the dtype's nesting\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\.conda\\envs\\runsit\\lib\\site-packages\\numpy\\lib\\npyio.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m    995\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    996\u001b[0m             \u001b[1;31m# Convert each value according to its column and store\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 997\u001b[1;33m             \u001b[0mitems\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mconv\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mval\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mconv\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mval\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mconverters\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mvals\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    998\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    999\u001b[0m             \u001b[1;31m# Then pack it according to the dtype's nesting\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\.conda\\envs\\runsit\\lib\\site-packages\\numpy\\lib\\npyio.py\u001b[0m in \u001b[0;36mfloatconv\u001b[1;34m(x)\u001b[0m\n\u001b[0;32m    732\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[1;34m'0x'\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    733\u001b[0m             \u001b[1;32mreturn\u001b[0m \u001b[0mfloat\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfromhex\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 734\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mfloat\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    735\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    736\u001b[0m     \u001b[0mtyp\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mdtype\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtype\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mValueError\u001b[0m: could not convert string to float: 'None'"
     ]
    }
   ],
   "source": [
    "if not trial > 0:\n",
    "    raise SystemExit('Nothing recorded. No sense to continue.')\n",
    "\n",
    "# do pack data to HDF5\n",
    "h5name = pack(session_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot sessions stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import h5py\n",
    "import numpy as np\n",
    "from scipy import signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arena_r = 0.4  # in meters\n",
    "\n",
    "with h5py.File(h5name, 'r') as f:\n",
    "    tl = np.array(f['processed']['timeline'])\n",
    "    trial_idxs = np.array(f['processed']['trial_idxs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12, 12))\n",
    "\n",
    "# trajectory and islands\n",
    "ax = fig.add_subplot(221)\n",
    "ax.scatter(tl[:, 1], tl[:, 2], s=1, alpha=0.1)  # positions\n",
    "scat = ax.scatter(trial_idxs[:, 2], trial_idxs[:, 3], s=1000, facecolors='none', edgecolors='r')  # islands, radius approx.\n",
    "ax.add_patch(plt.Circle((0, 0), arena_r, color='r', fill=False))\n",
    "ax.set_aspect('equal')\n",
    "ax.set_xlabel('X, m', fontsize=14)\n",
    "ax.set_ylabel('Y, m', fontsize=14)\n",
    "ax.set_title('Running', fontsize=14)\n",
    "ax.grid()\n",
    "\n",
    "# occupancy\n",
    "sigma = 0.1\n",
    "lin_profile = np.linspace(-15, 15, 20)\n",
    "bump = np.exp(-sigma * lin_profile**2)\n",
    "bump /= np.trapz(bump)  # normalize to 1\n",
    "kernel = bump[:, np.newaxis] * bump[np.newaxis, :]\n",
    "occupancy_map, _, _ = np.histogram2d(tl[:, 1], tl[:, 2], bins=[40, 40], range=np.array([[-0.5, 0.5], [-0.5, 0.5]]))\n",
    "occupancy_map = signal.convolve2d(occupancy_map, kernel, mode='same')\n",
    "\n",
    "ax = fig.add_subplot(222)\n",
    "ax.imshow(occupancy_map, origin='lower', extent=(-0.5, 0.5, -0.5, 0.5), cmap='Blues')\n",
    "ax.add_patch(plt.Circle((0, 0), arena_r, color='r', fill=False))\n",
    "ax.set_xlabel('X, m', fontsize=14)\n",
    "ax.set_title('Occupancy', fontsize=14)\n",
    "ax.grid()\n",
    "\n",
    "# trials\n",
    "durations = tl[trial_idxs[:, 1].astype(int)][:, 0] - tl[trial_idxs[:, 0].astype(int)][:, 0]\n",
    "colors = ['red' if x == 1 else 'grey' for x in trial_idxs[:, 5]]\n",
    "\n",
    "ax = fig.add_subplot(223)\n",
    "ax.barh(np.arange(len(trial_idxs)), durations, color=colors, align='center')\n",
    "ax.set_xlabel('Time, s', fontsize=14)\n",
    "ax.set_ylabel('Trial, #', fontsize=14)\n",
    "ax.set_title('Trials', fontsize=14)\n",
    "\n",
    "# speed\n",
    "ax = fig.add_subplot(224)\n",
    "ax.hist(tl[:, 3], bins=50, ec='black')\n",
    "ax.set_xlabel('Speed, m/s', fontsize=14)\n",
    "ax.set_title('Speed', fontsize=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
