{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['MKL_THREADING_LAYER'] = 'GNU'\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"16\"\n",
    "os.environ[\"MKL_NUM_THREADS\"] = \"16\"\n",
    "import torch\n",
    "import numpy as np\n",
    "import gymnasium as gym\n",
    "from gym_env import CustomEnv\n",
    "from stable_baselines3 import PPO\n",
    "from stable_baselines3.common.vec_env import SubprocVecEnv\n",
    "from stable_baselines3.common.vec_env import DummyVecEnv\n",
    "from sb3_contrib import RecurrentPPO\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from collections import deque\n",
    "\n",
    "\n",
    "\n",
    "current_dir = os.path.dirname(os.path.abspath(\"__file__\"))\n",
    "parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def make_env(gpu_id):\n",
    "#     def _init():\n",
    "#         os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)\n",
    "#         return CustomEnv(devicenum=gpu_id)\n",
    "#     return _init\n",
    "\n",
    "# env_fns = [make_env(i) for i in range(3)]\n",
    "# vec_env = SubprocVecEnv(env_fns)\n",
    "\n",
    "# vec_env = gym.make(\"CartPole-v1\", render_mode=\"human\")\n",
    "vec_env = CustomEnv(device_id=2)\n",
    "\n",
    "# model = RecurrentPPO(\n",
    "#     \"MlpLstmPolicy\", \n",
    "#     env=vec_env, \n",
    "#     # n_steps=320,\n",
    "#     # tensorboard_log=\"./tensorboard/\", \n",
    "#     # log_name=\"Lstm2\",\n",
    "#     device=torch.device(\"cuda:1\"),\n",
    "#     verbose=0)\n",
    "model = PPO.load(os.path.join(parent_dir, \"models\", \"d1a3o12_a2.zip\"), env=vec_env, device=\"cuda:2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_env = model.get_env()\n",
    "obs = test_env.reset()\n",
    "\n",
    "list_reward = []\n",
    "list_action = []\n",
    "list_cd = []\n",
    "list_cl = []\n",
    "list_sensor = []\n",
    "fifo_states = deque(maxlen=120)\n",
    "dones = [False for _ in range(test_env.num_envs)]\n",
    "# while not all(dones):\n",
    "for step in range(320):\n",
    "    action, _states = model.predict(observation=obs, deterministic=True)\n",
    "    # action = np.array([0.0, -1, +1], dtype=np.float32)[np.newaxis, :]\n",
    "    obs, rewards, dones, info = test_env.step(action)\n",
    "    list_reward.append(rewards)\n",
    "    list_action.append(action)\n",
    "    cd = [obs[0][0], obs[0][2], obs[0][4]]\n",
    "    cl = [obs[0][1], obs[0][3], obs[0][5]]\n",
    "    fifo_states.append(obs[0][6:12])\n",
    "    list_cd.append(cd)\n",
    "    list_cl.append(cl)\n",
    "    list_sensor.append(np.array(obs[0][6:12]))\n",
    "    # if step % 8 == 0:\n",
    "    #     test_env.envs[0].save(numstep=step)\n",
    "avg_reward = np.mean(list_reward[-160:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 524000 characters]",
      "text/plain": [
       "<Figure size 1200x1200 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# with open('Sin_PB_Best6.pkl', 'rb') as f:\n",
    "#     data = pickle.load(f)\n",
    "\n",
    "# 偏移量\n",
    "offset = np.array([0, -4, +4], dtype=np.float32)\n",
    "\n",
    "# 添加偏移量\n",
    "list_action_offset = [action[0]*8+offset for action in list_action]\n",
    "\n",
    "# Sum 3 CD and CL to get the total CD and CL\n",
    "list_overall_cd = [sum(cd) for cd in list_cd]\n",
    "list_overall_cl = [sum(cl) for cl in list_cl]\n",
    "\n",
    "n_points = 640\n",
    "plt.figure(figsize=(12, 12))\n",
    "plt.subplot(3, 2, 1)\n",
    "plt.plot(list_cd[:n_points])\n",
    "plt.title('CD over steps')\n",
    "plt.legend()\n",
    "plt.subplot(3, 2, 2)\n",
    "plt.plot(list_cl[:n_points])\n",
    "plt.title('CL over steps')\n",
    "plt.subplot(3, 2, 3)\n",
    "plt.plot(list_reward[:n_points])\n",
    "plt.title('Reward over steps')\n",
    "plt.subplot(3, 2, 4)\n",
    "plt.plot(list_overall_cd[:n_points], label='CD')\n",
    "plt.plot(list_overall_cl[:n_points], label='CL')\n",
    "plt.title('CD and CL over steps')\n",
    "plt.subplot(3, 2, 5)\n",
    "plt.plot(list_action_offset[:n_points])\n",
    "plt.title('Action over steps')\n",
    "plt.subplot(3, 2, 6)\n",
    "plt.plot(fifo_states)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pycuda_3_10",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
