{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2712d556-224f-4bf3-aced-6c50f2686abc",
   "metadata": {},
   "source": [
    "# Data Frames\n",
    "\n",
    "\n",
    "A `a2rl.WiDataFrame` (and `a2rl.Series`) is the central data structure leveraged by the rest of\n",
    "`a2rl` library such as simulators, tokenizers, and agents. It is a data frame that represents the\n",
    "historical *states*, *actions* and *rewards*, possibly with additional context. Each row in the\n",
    "dataframe denotes a state at a point-in-time, the action taken on that state (& time), and the\n",
    "rewards of taking that action on that specific state.\n",
    "\n",
    "A `a2rl.WiDataFrame` can be created from:\n",
    "\n",
    "- an existing data structure: `pandas.DataFrame`, ndarray (structured or homogeneous), Iterable,\n",
    "  or dict. See the `data` argument in\n",
    "  [pandas.DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html).\n",
    "- a `whatif` dataset, which is a directory comprising of a metadata YAML file, and a CSV data\n",
    "  file.\n",
    "\n",
    "**NOTE**: to enable color outputs in the notebook, please install\n",
    "[rich](https://github.com/Textualize/rich) (e.g., by running `%pip install rich` in your notebook)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95c6e9e4-7196-48a2-9e88-4242bd5200cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from __future__ import annotations\n",
    "import pandas as pd\n",
    "\n",
    "import a2rl as wi\n",
    "from a2rl.nbtools import pprint, print  # Enable color outputs when rich is installed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db46790a-0523-4069-9bbb-17407c597e58",
   "metadata": {},
   "source": [
    "## Load Dataset\n",
    "\n",
    "`whatif` comes with a few sample datasets to jumpstart your experience. You can list the available\n",
    "sample datasets, resolve the path of a specific sample dataset, and load a sample dataset into a\n",
    "`WiDataFrame`.\n",
    "\n",
    "**NOTE**: please see a separate example on how to create a `whatif` dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aa7ac1e-a4fd-4c84-9ede-500f5a335b53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List sample datasets\n",
    "print(\"Available sample datasets:\", wi.list_sample_datasets())\n",
    "\n",
    "# Get the path to the chiller dataset.\n",
    "p = wi.sample_dataset_path(\"chiller\")\n",
    "print(\"Location of chiller dataset:\", repr(p))\n",
    "\n",
    "# Load dataset to a dataframe\n",
    "df = wi.read_csv_dataset(p)\n",
    "pprint(\"Sar columns:\", df.sar_d)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68648aa6-d99b-46c5-afaa-7b84767be90b",
   "metadata": {},
   "source": [
    "You can see that the chiller dataframe specifies four *expected* sar columns, however the\n",
    "dataframe itself contains an additional column `timestamp`. This behavior is by-design, intended\n",
    "to allow you to apply additional preprocessing that require those additional contexts (i.e.,\n",
    "columns). In this particular chiller example, you might want to resample the dataframe by the\n",
    "`timestamp` column. Similarly, other operations can be applied.\n",
    "\n",
    "At some point in time, to train a simulator, you need a dataframe with strictly the *sar* columns\n",
    "only, which you can achieve by trimming the dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d1c7136-8649-477b-b457-9519b7e62bdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_trimmed = df.trim()\n",
    "df_trimmed.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f48055-8363-4511-90ed-206cff4dd33a",
   "metadata": {},
   "source": [
    "## Construct Data Frame\n",
    "\n",
    "You can also directly create a `WiDataFrame` from a plain `pandas.DataFrame`, or from any data\n",
    "source that you can use to create a regular `pandas.DataFrame`. However, `WiDataFrame` requires you\n",
    "to also supply the *sar* information.\n",
    "\n",
    "The following cell shows two typical variants to construct `WiDataFrame`: one from a dictionary, and\n",
    "another `pandas.DataFrame`. Please refer to\n",
    "[pandas.DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) to learn the\n",
    "additional types of data sources supported.\n",
    "\n",
    "You'll notice that the created data frames contains an additional columns, `z`, which is not part\n",
    "of *sar*. We will explain about this behavior in the next [section](#SAR-Columns)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f00fdff-747b-4652-97cc-10270f8483c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_from_dict = wi.WiDataFrame(\n",
    "    {\n",
    "        \"s1\": [1, 2, 3],\n",
    "        \"s2\": [3, 4, 5],\n",
    "        \"z\": [6, 7, 8],\n",
    "        \"a\": [\"x\", \"y\", \"z\"],\n",
    "        \"r\": [0.5, 1.5, 2.5],\n",
    "    },\n",
    "    states=[\"s1\", \"s2\"],\n",
    "    actions=[\"a\"],\n",
    "    rewards=[\"r\"],\n",
    ")\n",
    "\n",
    "display(df_from_dict.sar_d, df_from_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36bae583-fbd7-4026-becc-43718d60f6d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd_df = pd.DataFrame(\n",
    "    {\n",
    "        \"s1\": [10, 20, 30],\n",
    "        \"s2\": [30, 40, 50],\n",
    "        \"z\": [60, 70, 80],\n",
    "        \"a\": [\"X\", \"Y\", \"Z\"],\n",
    "        \"r\": [10.5, 11.5, 12.5],\n",
    "    },\n",
    ")\n",
    "\n",
    "df_from_pd_df = wi.WiDataFrame(\n",
    "    pd_df,\n",
    "    states=[\"s1\", \"s2\"],\n",
    "    actions=[\"a\"],\n",
    "    rewards=[\"r\"],\n",
    ")\n",
    "display(df_from_pd_df.sar_d, df_from_pd_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d5c34ed-71f5-4ada-8ede-641aa9810f1c",
   "metadata": {},
   "source": [
    "## SAR Information\n",
    "\n",
    "You can query the *sar* information as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fff05ee-0d3f-42a0-bf16-60ca1fecc62f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(\"Sar as dictionary:\", df.sar_d, \"\")\n",
    "pprint(\"Sar as list:\", df.sar, \"\")\n",
    "print(\"States:\", df.states)\n",
    "print(\"Actions:\", df.actions)\n",
    "print(\"Rewards:\", df.rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51f6fc3f-05e6-40de-90c1-30f34260a4c2",
   "metadata": {},
   "source": [
    "By now, you're already familiar with the key concept that `WiDataFrame` knows what should be its\n",
    "*sar* columns. However, it does not mean that those columns must really exist in the data frame.\n",
    "The dataframe itself may contains columns not in *sar* (e.g., the `timestamp` column in the chiller\n",
    "data), or even none of the *sar* columns at all!\n",
    "\n",
    "This behavior is by design. The intent is to let you specify only *once* the *sar* information of\n",
    "your business problem, and let `whatif` to manage and automatically propagate the association\n",
    "between your historical data and the *sar* information to the rest of `whatif` API.\n",
    "\n",
    "This design brings a few benefits:\n",
    "\n",
    "1. you don't need to manually keep track the association between a pandas data frame with your *sar*\n",
    "   information.\n",
    "2. you only need to pass around a `WiDataFrame` object instead of always passing both a\n",
    "   `pandas.DataFrame` and your *sar* information to various `whatif` operations.\n",
    "\n",
    "Let's illustrate these points with a simplified example which typically happens during\n",
    "preprocessing: splitting a dataframe into series, do something with the series, then combine the\n",
    "processed series back to a data frame again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87f24f07-36f6-4fa3-83ae-e21d974d6f82",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def assert_same_sar(sers: list[pd.Series]) -> None:\n",
    "    \"\"\"Check that all series has the same sar.\"\"\"\n",
    "    for ser in sers[1:]:\n",
    "        if ser.sar != sers[0].sar:\n",
    "            raise ValueError(\"Some series have different sar\")\n",
    "\n",
    "\n",
    "# Split the chiller df into 5 series (i.e., 1 for each column)\n",
    "if pd.__version__ >= \"1.5.0\":\n",
    "    sers = [ser for _, ser in df.items()]\n",
    "else:\n",
    "    sers = [ser for _, ser in df.iteritems()]\n",
    "assert_same_sar(sers)\n",
    "\n",
    "# Scale the states and rewards\n",
    "cols_to_scale = {*df.states, *df.rewards}\n",
    "sers = [ser / 25.0 if ser.name in cols_to_scale else ser for ser in sers]\n",
    "assert_same_sar(sers)\n",
    "\n",
    "# Show series names, and the common sar inherited from the source dataframe.\n",
    "pprint(\"All series have the same sar:\", sers[0].sar_d, \"\")\n",
    "\n",
    "# Reconstruct to a new df, and notice how the sar property is propagated.\n",
    "df_joined = pd.concat(sers, axis=1)\n",
    "pprint(\"Reconstructed dataframe has these sar:\", df_joined.sar_d)\n",
    "df_joined.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b866b39-a30c-4603-9051-708163835746",
   "metadata": {},
   "source": [
    "## Pandas Operations\n",
    "\n",
    "You can apply a wide range of pandas operations on `WiDataFrame` and `WiSeries`, just like to their\n",
    "`pandas` counterpart, and the results will still inherit the *sar* information. In fact, `pandas`\n",
    "operations applied to `whatif` data structures (i.e., frames and series) results in new `whatif`\n",
    "data structures.\n",
    "\n",
    "You've seen [slicing](https://pandas.pydata.org/docs/user_guide/indexing.html) and\n",
    "[concatenation](https://pandas.pydata.org/docs/user_guide/merging.html#concatenating-objects) in\n",
    "the previous cell. The next cells demonstrates a few more operations typically used during EDA and\n",
    "preprocessing. Notice how the *sar* information are propagated throughout the processing chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03fc3f2f-3a3e-47c3-a62a-ed3d477fccf5",
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "# Plot the states columns\n",
    "df[df.states].plot();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "935407fb-0b92-4e2c-8891-a26d9a045213",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Groupby: to compute average reward per action.\n",
    "df_mean_reward_per_action = df.groupby(df.actions).agg({df.rewards[0]: \"mean\"})\n",
    "pprint(\"Sar information:\", df_mean_reward_per_action.sar_d)\n",
    "df_mean_reward_per_action"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83db3315-217c-4543-8154-8508b9fa7855",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Congratulations! You've completed the tutorial on `whatif` data frame. We encourage you to further\n",
    "explore the remaining examples."
   ]
  }
 ],
 "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.6"
  },
  "toc-autonumbering": true,
  "vscode": {
   "interpreter": {
    "hash": "22f92e4608f34d3393fc5e7884f8906c6794e2d0198ea9b43992c442775a4328"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}