{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import pprint\n",
    "import numpy as np\n",
    "# # Load the .pkl file into a Python object\n",
    "# with open('/home/worker/yt/CVPRW23/Social-Implicit/checkpoint/social-implicit-hotel/args.pkl', 'rb') as f:\n",
    "#     data = pickle.load(f)\n",
    "\n",
    "# # Use pprint to pretty-print the object's structure\n",
    "# pprint.pprint(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "array([[ 0.        , -0.7400007 , -1.5800009 , -2.3700004 , -3.1400003 ,\n",
      "        -3.8400006 , -4.4500003 , -5.0700006 ],\n",
      "       [ 0.        ,  0.26999998,  0.37000036,  0.5300002 ,  0.6500001 ,\n",
      "         0.71000004,  0.8500004 ,  1.0100002 ]], dtype=float32)\n",
      "array([[ -5.4400005 ,  -5.8       ,  -6.1100006 ,  -6.3600006 ,\n",
      "         -6.84      ,  -7.4900007 ,  -8.3       ,  -9.030001  ,\n",
      "         -9.77      , -10.490001  , -11.14      , -11.83      ],\n",
      "       [  1.19      ,   1.6100001 ,   1.3300004 ,   1.7400002 ,\n",
      "          1.8900003 ,   2.0300002 ,   2.0300002 ,   1.8500004 ,\n",
      "          1.4300003 ,   1.0900002 ,   0.46000004,   0.0800004 ]],\n",
      "      dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# Load the .pkl file into a Python object\n",
    "with open('./Visualization/aligned/vis_eth.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)\n",
    "\n",
    "# Use pprint to pretty-print the object's structure\n",
    "pprint.pprint(data[0]['Observation'])\n",
    "pprint.pprint(data[0]['GroundTruth'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./Visualization/raw/Social-Implicit/save_eth.pkl', 'rb') as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['obs', 'tarj', 'pred'])\n"
     ]
    }
   ],
   "source": [
    "print(data.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['obs'] = np.concatenate(data['obs'], axis=1).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(181, 2, 8)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['obs'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "data['tarj'] = np.concatenate(data['tarj'], axis=1).squeeze()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(181, 2, 12)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['tarj'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "data['tarj'] = data['tarj'] - data['obs'][..., 0:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(181, 2, 12)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['tarj'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[10.31,  9.57,  8.73, ...,  6.47,  5.86,  5.24],\n",
       "        [ 5.97,  6.24,  6.34, ...,  6.68,  6.82,  6.98]],\n",
       "\n",
       "       [[12.49, 11.94, 11.03, ...,  8.59,  7.78,  6.96],\n",
       "        [ 6.6 ,  6.77,  6.84, ...,  6.85,  6.84,  6.84]],\n",
       "\n",
       "       [[12.51, 11.54, 10.96, ...,  9.54,  8.87,  8.04],\n",
       "        [ 6.19,  6.03,  5.97, ...,  6.09,  5.99,  5.66]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[ 0.87,  1.33,  1.79, ...,  3.24,  3.72,  4.14],\n",
       "        [ 7.36,  7.27,  7.27, ...,  7.11,  7.14,  7.14]],\n",
       "\n",
       "       [[ 1.33,  1.79,  2.25, ...,  3.68,  4.14,  4.57],\n",
       "        [ 6.6 ,  6.61,  6.52, ...,  6.54,  6.55,  6.56]],\n",
       "\n",
       "       [[ 1.33,  1.79,  2.28, ...,  3.72,  4.14,  4.61],\n",
       "        [ 7.27,  7.27,  7.21, ...,  7.14,  7.14,  7.15]]])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['obs']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[10.31],\n",
       "       [ 5.97]])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['obs'][..., 0:1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "PRED = []\n",
    "for i in range(len(data['pred'])):\n",
    "    PRED.append(np.stack(data['pred'][i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 12, 2, 2)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PRED[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "70"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(PRED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "PRED = []\n",
    "for i in range(len(data['pred'])):\n",
    "    PRED.append(np.stack(data['pred'][i]))\n",
    "data['pred'] = np.transpose(\n",
    "    np.concatenate(PRED, axis=2).squeeze(),\n",
    "    axes=(2, 0, 3, 1)) - data['obs'][:, None, :, 0:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(181, 100, 2, 12)"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['pred'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "data['obs'] = data['obs'] - data['obs'][..., 0:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.  , -0.74, -1.58, ..., -3.84, -4.45, -5.07],\n",
       "        [ 0.  ,  0.27,  0.37, ...,  0.71,  0.85,  1.01]],\n",
       "\n",
       "       [[ 0.  , -0.55, -1.46, ..., -3.9 , -4.71, -5.53],\n",
       "        [ 0.  ,  0.17,  0.24, ...,  0.25,  0.24,  0.24]],\n",
       "\n",
       "       [[ 0.  , -0.97, -1.55, ..., -2.97, -3.64, -4.47],\n",
       "        [ 0.  , -0.16, -0.22, ..., -0.1 , -0.2 , -0.53]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[ 0.  ,  0.46,  0.92, ...,  2.37,  2.85,  3.27],\n",
       "        [ 0.  , -0.09, -0.09, ..., -0.25, -0.22, -0.22]],\n",
       "\n",
       "       [[ 0.  ,  0.46,  0.92, ...,  2.35,  2.81,  3.24],\n",
       "        [ 0.  ,  0.01, -0.08, ..., -0.06, -0.05, -0.04]],\n",
       "\n",
       "       [[ 0.  ,  0.46,  0.95, ...,  2.39,  2.81,  3.28],\n",
       "        [ 0.  ,  0.  , -0.06, ..., -0.13, -0.13, -0.12]]])"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['obs']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}