{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np \n",
    "import pickle \n",
    "import matplotlib.pyplot as plt \n",
    "from collections import OrderedDict\n",
    "np.set_printoptions(suppress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dset = 'eth' # 'eth', 'zara1','hotel'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "364"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Trajectorn++\n",
    "trajpp = np.load('./raw/Trajectron++/' + dset + '_test_traj++.npz')\n",
    "trajpp_other = np.load('./raw/Trajectron++/' + dset + '_test_traj_gmm.npz')\n",
    "\n",
    "trajpp_ = {}\n",
    "trajpp_['obs'] = np.transpose(trajpp_other['obs'], axes=(0, 2, 1))\n",
    "trajpp_['tarj'] = np.transpose(trajpp['gt_abs'], axes=(0, 2, 1))\n",
    "trajpp_['tarj'] = trajpp_['tarj'] - trajpp_['obs'][..., 0:1]\n",
    "trajpp_['tarj_ref'] = np.transpose(trajpp['gt_rel'], axes=(0, 2, 1))\n",
    "trajpp_['pred'] = np.transpose(trajpp['pred_abs'], axes=(\n",
    "    0, 1, 3, 2))[:, :100, ...] - trajpp_['obs'][:, None, :, 0:1]\n",
    "trajpp_['obs'] = trajpp_['obs'] - trajpp_['obs'][..., 0:1]\n",
    "\n",
    "len(trajpp_['obs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "181"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./raw/Social-STGCNN/save_' + dset + '.pkl', 'rb') as f:\n",
    "    sstgcnn = pickle.load(f)\n",
    "\n",
    "sstgcnn_ = {}\n",
    "sstgcnn_['obs'] = np.concatenate(sstgcnn['obs'], axis=1).squeeze()\n",
    "sstgcnn_['tarj'] = np.concatenate(sstgcnn['tarj'], axis=1).squeeze()\n",
    "sstgcnn_['tarj'] = sstgcnn_['tarj'] - sstgcnn_['obs'][..., 0:1]\n",
    "full_traj = np.concatenate((sstgcnn_['obs'], sstgcnn_['tarj']), axis=-1)\n",
    "sstgcnn_['tarj_ref'] = (full_traj[..., 1:] - full_traj[..., :-1])[..., -12:]\n",
    "PRED = []\n",
    "for i in range(len(sstgcnn['pred'])):\n",
    "    PRED.append(np.stack(sstgcnn['pred'][i]))\n",
    "sstgcnn_['pred'] = np.transpose(\n",
    "    np.concatenate(PRED, axis=2).squeeze(),\n",
    "    axes=(2, 0, 3, 1)) - sstgcnn_['obs'][:, None, :, 0:1]\n",
    "sstgcnn_['obs'] = sstgcnn_['obs'] - sstgcnn_['obs'][..., 0:1]\n",
    "len(sstgcnn_['obs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "181"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./raw/Social-Implicit/save_' + dset + '.pkl', 'rb') as f:\n",
    "    simplicit = pickle.load(f)\n",
    "\n",
    "simplicit_ = {}\n",
    "simplicit_['obs'] = np.concatenate(simplicit['obs'], axis=1).squeeze()\n",
    "simplicit_['tarj'] = np.concatenate(simplicit['tarj'], axis=1).squeeze()\n",
    "simplicit_['tarj'] = simplicit_['tarj'] - simplicit_['obs'][..., 0:1]\n",
    "full_traj = np.concatenate((simplicit_['obs'], simplicit_['tarj']), axis=-1)\n",
    "simplicit_['tarj_ref'] = (full_traj[..., 1:] - full_traj[..., :-1])[..., -12:]\n",
    "PRED = []\n",
    "for i in range(len(simplicit['pred'])):\n",
    "    PRED.append(np.stack(simplicit['pred'][i]))\n",
    "simplicit_['pred'] = np.transpose(\n",
    "    np.concatenate(PRED, axis=2).squeeze(),\n",
    "    axes=(2, 0, 3, 1)) - simplicit_['obs'][:, None, :, 0:1]\n",
    "simplicit_['obs'] = simplicit_['obs'] - simplicit_['obs'][..., 0:1]\n",
    "\n",
    "len(simplicit_['obs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "364"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./raw/ExpertTraj/save_' + dset + '.pkl', 'rb') as f:\n",
    "    expertraj = pickle.load(f)\n",
    "\n",
    "expertraj_ = {}\n",
    "expertraj_['obs'] = np.transpose(np.concatenate(expertraj['obs'],\n",
    "                                                axis=1).squeeze(),\n",
    "                                 axes=(1, 2, 0))\n",
    "expertraj_['tarj'] = np.transpose(np.concatenate(expertraj['tarj'],\n",
    "                                                 axis=1).squeeze(),\n",
    "                                  axes=(1, 2, 0))\n",
    "full_traj = np.concatenate((expertraj_['obs'], expertraj_['tarj']), axis=-1)\n",
    "expertraj_['tarj_ref'] = (full_traj[..., 1:] - full_traj[..., :-1])[..., -12:]\n",
    "PRED = []\n",
    "for i in range(len(expertraj['pred'])):\n",
    "    PRED.append(np.stack(expertraj['pred'][i]))\n",
    "expertraj_['pred'] = np.transpose(np.concatenate(PRED, axis=2).squeeze(),\n",
    "                                  axes=(2, 0, 3, 1))\n",
    "\n",
    "len(expertraj_['obs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "181"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('./raw/Social-GAN/save_'+dset+'.pkl','rb') as f:\n",
    "    sgan = pickle.load(f)\n",
    "\n",
    "sgan_ = {}\n",
    "sgan_['obs'] = np.transpose(np.concatenate(sgan['obs'],axis = 1).squeeze(), axes=(1,2,0))\n",
    "sgan_['tarj'] = np.transpose(np.concatenate(sgan['tarj'],axis = 1).squeeze(), axes=(1,2,0))\n",
    "sgan_['tarj'] = sgan_['tarj']-sgan_['obs'][...,0:1]\n",
    "full_traj =np.concatenate((sgan_['obs'], sgan_['tarj']),axis=-1)\n",
    "sgan_['tarj_ref'] = (full_traj[...,1:]- full_traj[...,:-1])[...,-12:]\n",
    "PRED = []\n",
    "for i in range(len(sgan['pred'])):\n",
    "    PRED.append(np.stack(sgan['pred'][i]))\n",
    "sgan_['pred'] = np.transpose(np.concatenate(PRED,axis = 2).squeeze(), axes=(2,0,3,1))- sgan_['obs'][:, None, :, 0:1]\n",
    "sgan_['obs'] = sgan_['obs']- sgan_['obs'][...,0:1]\n",
    "\n",
    "len(sgan_['tarj'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_closer_index(key,query):\n",
    "    diff = np.linalg.norm(query - key, axis=-2).sum(-1)\n",
    "    argmin = np.argmin(diff,axis=0)\n",
    "    # print(diff[argmin], argmin, \"\\n\",key,\"\\n\",query[argmin])\n",
    "    return argmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "collect = []\n",
    "\n",
    "for key in  range(len(sstgcnn_['obs'])):\n",
    "    collect.append(OrderedDict())\n",
    "    key_value = sstgcnn_['obs'][key]\n",
    "\n",
    "    key_trajpp = get_closer_index(key_value, trajpp_['obs'])\n",
    "    key_sgan = get_closer_index(key_value, sgan_['obs'])\n",
    "    key_simplicit = get_closer_index(key_value, simplicit_['obs'])\n",
    "    key_expertraj = get_closer_index(key_value, expertraj_['obs'])\n",
    "\n",
    "    collect[-1]['Observation'] = sstgcnn_['obs'][key]\n",
    "    collect[-1]['GroundTruth'] = sstgcnn_['tarj'][key]\n",
    "    collect[-1]['S-GAN'] = sgan_['pred'][key_sgan]\n",
    "    collect[-1]['S-STGCNN'] = sstgcnn_['pred'][key]\n",
    "    collect[-1]['S-Implicit'] = simplicit_['pred'][key_simplicit]\n",
    "    collect[-1]['Trajectron++'] = trajpp_['pred'][key_trajpp]\n",
    "    collect[-1]['ExpertTraj'] = expertraj_['pred'][key_expertraj]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "181"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# print stats \n",
    "len(collect)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('./aligned'):\n",
    "    os.makedirs('./aligned')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./aligned/vis_'+dset+'.pkl','wb') as f : \n",
    "    pickle.dump(collect,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}