{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "from stable_baselines3 import PPO\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)\n",
    "\n",
    "cuda.init()\n",
    "context2 = cuda.Device(2).make_context()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(output_dir, \"d1a3o12_re100.pkl\"), 'rb') as f:\n",
    "    data = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gym_env import CustomEnv\n",
    "context2.push()\n",
    "vec_env = CustomEnv(device_id=2)\n",
    "context2.pop()\n",
    "model = PPO.load(os.path.join(parent_dir, \"models\", \"d1a3o12_re100.zip\"), env=vec_env, device=\"cuda:2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context2.push()\n",
    "model_env = model.get_env()\n",
    "# obs = model_env.reset()\n",
    "obs_init = np.array(model_env.envs[0].save_states.copy())\n",
    "context2.pop()\n",
    "\n",
    "obs_init[:, 6:12] = obs_init[:, 6:12]/model_env.envs[0].force_norm_fact.copy()\n",
    "sens_deviation = model_env.envs[0].sens_deviation.copy()\n",
    "sens_norm_fact = model_env.envs[0].sens_norm_fact.copy()\n",
    "for i in range(6):\n",
    "    obs_init[:, i] = (obs_init[:, i] - sens_deviation[i]) / sens_norm_fact[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from darts import TimeSeries\n",
    "from darts.models import RNNModel  # 或者 NBEATSModel、TFTModel 等\n",
    "from darts.models import TFTModel\n",
    "from darts.dataprocessing.transformers import Scaler\n",
    "from typing import List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_obs_train = []\n",
    "set_act_train = []\n",
    "index = pd.RangeIndex(start=0, stop=480, step=1)\n",
    "for i in range(len(data)):\n",
    "    obs_episolde = np.array(data[i]['observations'])[:, 0]\n",
    "    if np.any(np.all(obs_episolde == 0, axis=1)):\n",
    "        # print(f\"Episode {i} contains a row with all zeros in observations.\")\n",
    "        continue\n",
    "    else:\n",
    "        obs_full = np.concatenate((obs_init, obs_episolde), axis=0)\n",
    "        act_init = np.zeros((len(obs_init), 3))\n",
    "        act_episolde = np.array(data[i]['actions'])\n",
    "        act_full = np.concatenate((act_init, act_episolde), axis=0)\n",
    "\n",
    "        obs_ts = TimeSeries.from_times_and_values(\n",
    "            times=index,\n",
    "            values=obs_full.astype(np.float32), \n",
    "            columns=[f\"obs{dim+1}\" for dim in range(12)]\n",
    "        )\n",
    "        act_ts = TimeSeries.from_times_and_values(\n",
    "            times=index,\n",
    "            values=act_full.astype(np.float32), \n",
    "            columns=[f\"act{dim+1}\" for dim in range(3)]\n",
    "        )\n",
    "\n",
    "        set_obs_train.append(obs_ts)\n",
    "        set_act_train.append(act_ts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Trainer will use only 1 of 4 GPUs because it is running inside an interactive / notebook environment. You may try to set `Trainer(devices=4)` but please note that multi-GPU inside interactive / notebook environments is considered experimental and unstable. Your mileage may vary.\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "   | Name                              | Type                             | Params | Mode \n",
      "------------------------------------------------------------------------------------------------\n",
      "0  | train_metrics                     | MetricCollection                 | 0      | train\n",
      "1  | val_metrics                       | MetricCollection                 | 0      | train\n",
      "2  | input_embeddings                  | _MultiEmbedding                  | 0      | train\n",
      "3  | static_covariates_vsn             | _VariableSelectionNetwork        | 0      | train\n",
      "4  | encoder_vsn                       | _VariableSelectionNetwork        | 17.8 K | train\n",
      "5  | decoder_vsn                       | _VariableSelectionNetwork        | 896    | train\n",
      "6  | static_context_grn                | _GatedResidualNetwork            | 4.3 K  | train\n",
      "7  | static_context_hidden_encoder_grn | _GatedResidualNetwork            | 4.3 K  | train\n",
      "8  | static_context_cell_encoder_grn   | _GatedResidualNetwork            | 4.3 K  | train\n",
      "9  | static_context_enrichment         | _GatedResidualNetwork            | 4.3 K  | train\n",
      "10 | lstm_encoder                      | LSTM                             | 8.4 K  | train\n",
      "11 | lstm_decoder                      | LSTM                             | 8.4 K  | train\n",
      "12 | post_lstm_gan                     | _GateAddNorm                     | 2.2 K  | train\n",
      "13 | static_enrichment_grn             | _GatedResidualNetwork            | 5.3 K  | train\n",
      "14 | multihead_attn                    | _InterpretableMultiHeadAttention | 2.6 K  | train\n",
      "15 | post_attn_gan                     | _GateAddNorm                     | 2.2 K  | train\n",
      "16 | feed_forward_block                | _GatedResidualNetwork            | 4.3 K  | train\n",
      "17 | pre_output_gan                    | _GateAddNorm                     | 2.2 K  | train\n",
      "18 | output_layer                      | Linear                           | 6.7 K  | train\n",
      "------------------------------------------------------------------------------------------------\n",
      "78.2 K    Trainable params\n",
      "0         Non-trainable params\n",
      "78.2 K    Total params\n",
      "0.313     Total estimated model params size (MB)\n",
      "382       Modules in train mode\n",
      "0         Modules in eval mode\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 49: 100%|██████████| 12/12 [00:00<00:00, 21.99it/s, train_loss=0.133]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=50` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 49: 100%|██████████| 12/12 [00:00<00:00, 21.93it/s, train_loss=0.133]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TFTModel(output_chunk_shift=0, hidden_size=32, lstm_layers=1, num_attention_heads=4, full_attention=False, feed_forward=GatedResidualNetwork, dropout=0.1, hidden_continuous_size=8, categorical_embedding_sizes=None, add_relative_index=True, loss_fn=None, likelihood=None, norm_type=LayerNorm, use_static_covariates=True, input_chunk_length=120, output_chunk_length=1, batch_size=32, n_epochs=50, random_state=42)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_tft = TFTModel(\n",
    "    input_chunk_length=120,\n",
    "    output_chunk_length=1,  # 预测1步\n",
    "    hidden_size=32,\n",
    "    lstm_layers=1,\n",
    "    dropout=0.1,\n",
    "    batch_size=32,\n",
    "    n_epochs=50,\n",
    "    add_relative_index=True,\n",
    "    random_state=42,\n",
    "    # 注意：TFT支持past_covariates与future_covariates，也可设置transformer头等更多超参\n",
    ")\n",
    "\n",
    "model_tft.fit(\n",
    "    series=obs_ts,\n",
    "    past_covariates=act_ts\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(360):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data = obs_full[:, 11]\n",
    "import matplotlib.pyplot as plt\n",
    "plt.plot(plot_data)"
   ]
  }
 ],
 "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
}
