{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cad57710",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91d021d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SITDisplay:\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), 'pink': (255, 0, 255), 'tortoise': (0, 255, 255)\n",
    "    }\n",
    "\n",
    "    island_colors = {1: 'pink', 2: 'green', 3: 'blue', 4: 'yellow', 5: 'black'}\n",
    "    animal_colors = {1: 'tortoise', 2: 'pink', 3: 'white'}\n",
    "\n",
    "    def __init__(self, pt, cfg):  # injected dependency - position controller\n",
    "        self.pt = pt\n",
    "        self.cfg = cfg\n",
    "        self.masked_frame = None\n",
    "        \n",
    "    def render(self, input_frame, status, islands=[], angles=[], text_infos=[]):\n",
    "        COLORS = SITDisplay.COLORS\n",
    "        FONT = cv2.FONT_HERSHEY_DUPLEX\n",
    "        status_color = COLORS['green'] if status.value == 1 else COLORS['red']\n",
    "        texts = list(text_infos)\n",
    "        \n",
    "        # mask unused areas\n",
    "        frame = cv2.bitwise_and(src1=input_frame, src2=self.pt.mask)\n",
    "        self.masked_frame = frame.copy()\n",
    "\n",
    "        # size of the arena and status indicator\n",
    "        cv2.circle(frame, (self.pt.cfg['arena_x'], self.pt.cfg['arena_y']), self.pt.cfg['arena_radius'], COLORS['red'], 2)\n",
    "        cv2.circle(frame, (self.pt.cfg['arena_x'], self.pt.cfg['arena_y']), self.pt.cfg['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[SITDisplay.island_colors[island.sound_id]]\n",
    "                x, y = self.pt.meters_to_px(island.x, island.y)\n",
    "                cv2.circle(frame, (x, y), int(island.r/self.pt.pixel_size), clr, 2)\n",
    "            \n",
    "        # draw anglular indicators\n",
    "        if len(angles) > 0:  # should not be more than 4\n",
    "            p_rho = self.pt.cfg['floor_radius'] * self.pt.pixel_size\n",
    "            for p_phi, clr in zip(angles, ['blue', 'yellow', 'pink', 'tortoise'][:len(angles)]):\n",
    "                x_init, y_init = self.pt.meters_to_px(p_rho * np.sin(p_phi), p_rho * np.cos(p_phi))\n",
    "                cv2.circle(frame, (x_init, y_init), 10, COLORS[clr], -1)\n",
    "\n",
    "        # light or dark period\n",
    "        texts.append('Light' if self.pt.is_light else 'Dark')\n",
    "\n",
    "        if not self.cfg['save_contours']:\n",
    "            frame_to_save = frame.copy()  # video without animal contours\n",
    "        \n",
    "        # draw animal position center(s) and contours\n",
    "        if self.pt.positions_in_px is not None:\n",
    "            pos_in_px = self.pt.positions_in_px\n",
    "            pos_in_m = self.pt.positions_in_m\n",
    "            contours = self.pt.contours\n",
    "            \n",
    "            for i in range(len(pos_in_px)):\n",
    "                clr = COLORS[SITDisplay.animal_colors[i+1]]\n",
    "                x_px, y_px, alpha = pos_in_px[i][0], pos_in_px[i][1], self.pt.hds[0]\n",
    "                \n",
    "                # draw animal center and contours\n",
    "                cv2.circle(frame, (x_px, y_px), 5, clr, -1)\n",
    "                cv2.drawContours(frame, [contours[i]], 0, clr, 1, cv2.LINE_AA)\n",
    "                \n",
    "                # draw (instant position values and) HD vector\n",
    "                #cv2.putText(frame, '%.2f %.2f' % (pos_in_m[i][0], pos_in_m[i][1]), (x_px + 10, y_px + 10), FONT, .3, clr)\n",
    "                #x1, y1 = int(x_px + 40*np.cos(np.deg2rad(alpha))), int(y_px - 40*np.sin(np.deg2rad(alpha)))\n",
    "                #cv2.line(frame, (x_px, y_px), (x1, y1), clr, 1)\n",
    "                \n",
    "                texts.append('Animal %d: %.3f %.3f' % (i+1, pos_in_m[i][0], pos_in_m[i][1]))\n",
    "                texts.append('Speed/HD: %.2f %.2f' % (self.pt.speeds[0], self.pt.hds[0]))\n",
    "\n",
    "        if self.cfg['save_contours']:\n",
    "            frame_to_save = frame.copy()  # save video with animal contours\n",
    "            \n",
    "        # draw text infos\n",
    "        for i, text in enumerate(texts):\n",
    "            cv2.putText(frame, text, (10, 30 + 20*(i+1)), FONT, .5, COLORS['white'])\n",
    "            cv2.putText(frame_to_save, text, (10, 30 + 20*(i+1)), FONT, .5, COLORS['white'])\n",
    "\n",
    "        return frame, frame_to_save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e7f5cbd",
   "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": 5
}