{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp data.utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Utils\n",
    "> Collection of utilities to work with Time Series data and tests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def create_synthetic_tsdata(n_ts: int = 64, \n",
    "                            sort: bool = False) -> Tuple[pd.DataFrame,\n",
    "                                                         pd.DataFrame,\n",
    "                                                         pd.DataFrame]:\n",
    "    \"\"\"Creates synthetic time serie data.\"\"\"\n",
    "    uids = np.array([f'uid_{i + 1}' for i in range(n_ts)])\n",
    "    dss = pd.date_range(end='2020-12-31', periods=n_ts)\n",
    "    \n",
    "    df = []\n",
    "    for idx in range(n_ts):\n",
    "        ts = pd.DataFrame({'unique_id': np.repeat(uids[idx], idx + 1),\n",
    "                           'ds': dss[-(idx + 1):],\n",
    "                           'y': 1 + np.arange(idx + 1)})\n",
    "        df.append(ts)\n",
    "    \n",
    "    df = pd.concat(df)\n",
    "    df['day_of_week'] = df['ds'].dt.day_of_week\n",
    "    df['future_1'] = df['y'] + 1\n",
    "    df['id_ts'] = df['unique_id'].astype('category').cat.codes\n",
    "    if sort:\n",
    "        df = df.sort_values(['unique_id', 'ds'])\n",
    "    \n",
    "    Y_df = df.filter(items=['unique_id', 'ds', 'y'])\n",
    "    X_df = df.filter(items=['unique_id', 'ds', 'day_of_week', 'future_1'])\n",
    "    S_df = df.filter(items=['unique_id', 'id_ts']).drop_duplicates()\n",
    "    \n",
    "    return Y_df, X_df, S_df"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nixtla",
   "language": "python",
   "name": "nixtla"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
