{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "from extract_consecutive_screenshots import get_pickle_name, load_objects, load_snapshot_data, analyze_run\n",
    "from visualization import get_connection, get_gta_image_jpg_dir\n",
    "from os import path\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from matplotlib import patches\n",
    "import numpy as np\n",
    "import progressbar\n",
    "import os\n",
    "# JPEG_DIR = get_gta_image_jpg_dir()\n",
    "JPEG_DIR = 'D:/GTAV_extraction_output/rgb-jpeg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "run_id = 52\n",
    "analyze_run(run_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6782"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cars = load_objects(run_id)\n",
    "len(cars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_pos_diffs(snapshots):\n",
    "    pos_diffs = []\n",
    "    positions = [snap['position'] for snap in snaps]\n",
    "    for i in range(1, len(snaps)):\n",
    "        diff = np.array(positions[i]) - np.array(positions[i - 1])\n",
    "        pos_diffs.append(diff)\n",
    "    return pos_diffs\n",
    "\n",
    "def get_time_diffs(snapshots):\n",
    "    # gets difference in seconds, for normalization\n",
    "    time_diffs = []\n",
    "    times = [snap['timestamp'] for snap in snapshots]\n",
    "    for i in range(1, len(snapshots)):\n",
    "        diff = (times[i] - times[i - 1]).total_seconds()\n",
    "        time_diffs.append(diff)\n",
    "    return time_diffs\n",
    "\n",
    "def transform_to_relative(diff, cam):\n",
    "    base = np.array([\n",
    "        [cam[0], cam[1]],\n",
    "        [cam[1], -cam[0]]\n",
    "    ])\n",
    "    return np.matmul(base, diff[0:2])\n",
    "\n",
    "def normalize_diff(relative_diff, time_diff):\n",
    "    return relative_diff / time_diff\n",
    "    \n",
    "def get_pickle_processed_name(run_id):\n",
    "    return os.path.join('runs', 'run_{}_processed_pickle.rick'.format(run_id))\n",
    "\n",
    "def save_processed_objects(run_id, objects):\n",
    "    data_file = get_pickle_processed_name(run_id)\n",
    "    with open(data_file, 'wb+') as file:\n",
    "        pickle.dump(objects, file)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% 6782 |#######################################################|   2.0 KiB/s\n"
     ]
    }
   ],
   "source": [
    "widgets = [progressbar.Percentage(), ' ', progressbar.Counter(), ' ', progressbar.Bar(), ' ', progressbar.FileTransferSpeed()]\n",
    "pbar = progressbar.ProgressBar(widgets=widgets, max_value = len(cars)).start()\n",
    "counter = 0\n",
    "\n",
    "for i, tupl in enumerate(cars.items()):\n",
    "    key, car = tupl\n",
    "    pbar.update(i)\n",
    "    snaps = car['snapshots']\n",
    "    pos_diffs = get_pos_diffs(snaps)\n",
    "    time_diffs = get_time_diffs(snaps)\n",
    "    for i in range(len(pos_diffs)):\n",
    "        snap = snaps[i]\n",
    "        diff = pos_diffs[i]\n",
    "        time_diff = time_diffs[i]\n",
    "        diff_rel = transform_to_relative(diff, snap['cam_direction'][0:2])\n",
    "        diff_normalized = normalize_diff(diff_rel, time_diff)\n",
    "        \n",
    "        snap['diff'] = diff\n",
    "        snap['diff_relative'] = diff_rel\n",
    "        snap['diff_normalized'] = diff_normalized\n",
    "        snaps[i]\n",
    "pbar.finish()\n",
    "\n",
    "save_processed_objects(run_id, cars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
