{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle\n",
    "import pycuda.driver as cuda\n",
    "import sys\n",
    "import os\n",
    "from datetime import datetime\n",
    "\n",
    "current_dir = os.path.dirname(os.path.abspath(\"__file__\"))\n",
    "parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))\n",
    "output_dir = os.path.join(parent_dir, \"output\")\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "sys.path.append(parent_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "cuda.init()\n",
    "\n",
    "context = cuda.Device(3).make_context()\n",
    "\n",
    "DATA_TYPE = np.float32\n",
    "\n",
    "from env_manifold import CustomEnv\n",
    "context.push()\n",
    "env = CustomEnv(device_id=3)\n",
    "context.pop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = [2, 2, 2]\n",
    "\n",
    "def generate_random_group(size, low, high):\n",
    "    intervals = np.linspace(low, high, size + 1)\n",
    "    group = np.concatenate([np.random.uniform(intervals[i], intervals[i+1], 1) for i in range(size)])\n",
    "    return np.sort(group)\n",
    "\n",
    "group1 = generate_random_group(size[0], -1, 1)\n",
    "group2 = generate_random_group(size[1], -1, 1)\n",
    "group3 = generate_random_group(size[2], -1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-10-29 11:23:57 - (0, 0, 0)\n",
      "2024-10-29 11:24:39 - (0, 0, 1)\n",
      "2024-10-29 11:25:21 - (0, 1, 0)\n",
      "2024-10-29 11:26:03 - (0, 1, 1)\n",
      "2024-10-29 11:26:45 - (1, 0, 0)\n",
      "2024-10-29 11:27:27 - (1, 0, 1)\n",
      "2024-10-29 11:28:09 - (1, 1, 0)\n",
      "2024-10-29 11:28:51 - (1, 1, 1)\n"
     ]
    }
   ],
   "source": [
    "data = np.empty(size, dtype=object)\n",
    "for i, a1 in enumerate(sorted(group1)):\n",
    "    for j, a2 in enumerate(sorted(group2)):\n",
    "        for k, a3 in enumerate(sorted(group3)):\n",
    "            context.push()\n",
    "            action = np.array([a1, a2, a3], dtype=np.float32)\n",
    "            env.reset()\n",
    "            for _ in range(400):\n",
    "                _, _, _, _, _ = env.step(action)\n",
    "            context.pop()\n",
    "            fifo = np.array(env.fifo_states.copy())\n",
    "            data[i, j, k] = {'action': action, 'fifo': fifo}\n",
    "            current_time = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
    "            print(f\"{current_time} - ({i}, {j}, {k})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(output_dir, \"manifold_test.pkl\"), 'wb') as f:\n",
    "    pickle.dump(data, f)\n",
    "\n",
    "# with open(os.path.join(output_dir, \"manifold_test.pkl\"), 'rb') as f:\n",
    "#     data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.44718933 0.01895384 0.4848376 ]\n",
      "[ 2.3110747 -1.8382945  1.3480661  1.072888  -3.6636636 -4.1307545]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 30008 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_data = data[0,1,1]['fifo'][:,0]\n",
    "import matplotlib.pyplot as plt\n",
    "plt.plot(plot_data)\n",
    "print(data[1,1,1]['action'])\n",
    "print(np.mean(plot_data, axis=0))"
   ]
  }
 ],
 "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
}
