{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp data.tsdataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Base Dataset\n",
    "> Transforms pandas DataFrame into a TimeSeriesDataset for the Dataloder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import gc\n",
    "import logging\n",
    "from typing import Dict, List, Optional, Tuple, Union\n",
    "from typing_extensions import Literal\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch as t\n",
    "from fastcore.foundation import patch\n",
    "from torch.utils.data import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class BaseDataset(Dataset):\n",
    "    \"\"\"\n",
    "    A class used to store Time Series data.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,\n",
    "                 Y_df: pd.DataFrame,\n",
    "                 X_df: Optional[pd.DataFrame] = None,\n",
    "                 S_df: Optional[pd.DataFrame] = None,\n",
    "                 f_cols: Optional[List] = None,\n",
    "                 mask_df: Optional[pd.DataFrame] = None,\n",
    "                 ds_in_test: int = 0,\n",
    "                 is_test: bool = False,\n",
    "                 input_size: int = None,\n",
    "                 output_size: int = None,\n",
    "                 complete_windows: bool = True,\n",
    "                 verbose: bool = False) -> 'BaseDataset':\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        Y_df: pd.DataFrame\n",
    "            Target time series with columns ['unique_id', 'ds', 'y'].\n",
    "        X_df: pd.DataFrame\n",
    "            Exogenous time series with columns ['unique_id', 'ds', 'y'].\n",
    "        S_df: pd.DataFrame\n",
    "            Static exogenous variables with columns ['unique_id', 'ds'] \n",
    "            and static variables.\n",
    "        f_cols: list\n",
    "            List of exogenous variables of the future.\n",
    "        mask_df: pd.DataFrame\n",
    "            Outsample mask with columns ['unique_id', 'ds', 'sample_mask']\n",
    "            and optionally 'available_mask'.\n",
    "            Default None: constructs default mask based on ds_in_test.\n",
    "        ds_in_test: int\n",
    "            Only used when mask_df = None.\n",
    "            Numer of datestamps to use as outsample.\n",
    "        is_test: bool\n",
    "            Only used when mask_df = None.\n",
    "            Wheter target time series belongs to test set.\n",
    "        input_size: int\n",
    "            Size of the training sets.\n",
    "        output_size: int\n",
    "            Forecast horizon.\n",
    "        complete_windows: bool\n",
    "            Whether consider only windows with available window_size.\n",
    "            Default False.\n",
    "        verbose: bool\n",
    "            Wheter or not log outputs.\n",
    "        \"\"\"        \n",
    "        assert type(Y_df) == pd.core.frame.DataFrame\n",
    "        assert all([(col in Y_df) for col in ['unique_id', 'ds', 'y']])\n",
    "        self.verbose = verbose\n",
    "\n",
    "        if X_df is not None:\n",
    "            assert type(X_df) == pd.core.frame.DataFrame\n",
    "            assert all([(col in X_df) for col in ['unique_id', 'ds']])\n",
    "            assert len(Y_df)==len(X_df), 'The dimensions of Y_df and X_df are not the same'\n",
    "\n",
    "        if mask_df is not None:\n",
    "            assert len(Y_df)==len(mask_df), 'The dimensions of Y_df and mask_df are not the same'\n",
    "            assert all([(col in mask_df) for col in ['unique_id', 'ds', 'sample_mask']])\n",
    "            if 'available_mask' not in mask_df.columns:\n",
    "                if self.verbose: \n",
    "                    logging.info('Available mask not provided, defaulted with 1s.')\n",
    "                mask_df['available_mask'] = 1\n",
    "            assert np.sum(np.isnan(mask_df.available_mask.values)) == 0\n",
    "            assert np.sum(np.isnan(mask_df.sample_mask.values)) == 0\n",
    "        else:\n",
    "            mask_df = get_default_mask_df(Y_df=Y_df, \n",
    "                                          is_test=is_test,\n",
    "                                          ds_in_test=ds_in_test)\n",
    "        \n",
    "        n_ds  = len(mask_df)\n",
    "        n_avl = mask_df.available_mask.sum()        \n",
    "        n_ins = mask_df.sample_mask.sum()\n",
    "        n_out = len(mask_df) - mask_df.sample_mask.sum()\n",
    "\n",
    "        avl_prc = np.round((100 * n_avl) / n_ds, 2)\n",
    "        ins_prc = np.round((100 * n_ins) / n_ds, 2)\n",
    "        out_prc = np.round((100 * n_out) / n_ds, 2)\n",
    "        if self.verbose:\n",
    "            logging.info('Train Validation splits\\n')\n",
    "            if len(mask_df.unique_id.unique()) < 10:\n",
    "                logging.info(mask_df.groupby(['unique_id', 'sample_mask']).agg({'ds': ['min', 'max']}))\n",
    "            else:\n",
    "                logging.info(mask_df.groupby(['sample_mask']).agg({'ds': ['min', 'max']}))\n",
    "            dataset_info  = f'\\nTotal data \\t\\t\\t{n_ds} time stamps \\n'\n",
    "            dataset_info += f'Available percentage={avl_prc}, \\t{n_avl} time stamps \\n'\n",
    "            dataset_info += f'Insample  percentage={ins_prc}, \\t{n_ins} time stamps \\n'\n",
    "            dataset_info += f'Outsample percentage={out_prc}, \\t{n_out} time stamps \\n'\n",
    "            logging.info(dataset_info)\n",
    " \n",
    "        self.ts_data, self.s_matrix, self.meta_data, self.t_cols, self.s_cols \\\n",
    "                         = self._df_to_lists(Y_df=Y_df, S_df=S_df, X_df=X_df, mask_df=mask_df)\n",
    "\n",
    "        # Dataset attributes\n",
    "        self.n_series = len(self.ts_data)\n",
    "        self.max_len = max([len(ts) for ts in self.ts_data])\n",
    "        self.n_channels = len(self.t_cols) # t_cols insample_mask and outsample_mask\n",
    "        self.frequency = pd.infer_freq(Y_df.head()['ds'])\n",
    "        self.f_cols = f_cols\n",
    "        self.f_idxs = self._get_f_idxs(f_cols) if f_cols else []\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.complete_windows = complete_windows\n",
    "        self.first_ds = 0\n",
    "\n",
    "        # Number of X and S features\n",
    "        self.n_x = 0 if X_df is None else X_df.shape[1] - 2 # -2 for unique_id and ds\n",
    "        self.n_s = 0 if S_df is None else S_df.shape[1] - 1 # -1 for unique_id\n",
    "\n",
    "        # Balances panel and creates \n",
    "        # numpy  s_matrix of shape (n_series, n_s)\n",
    "        # numpy ts_tensor of shape (n_series, n_channels, max_len) n_channels = t_cols + masks\n",
    "        self.len_series, self.ts_tensor = self._create_tensor()\n",
    "        \n",
    "        # Defining sampleable time series\n",
    "        self.ts_idxs = np.arange(self.n_series)\n",
    "        self.sampleable_ts_idxs: np.ndarray\n",
    "        self.n_sampleable_ts: int\n",
    "            \n",
    "        self._define_sampleable_ts_idxs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def _define_sampleable_ts_idxs(self: BaseDataset):\n",
    "    self.n_sampleable_ts = len(self.ts_tensor)\n",
    "    self.sampleable_ts_idxs = self.ts_idxs.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def _df_to_lists(self: BaseDataset, \n",
    "                 S_df: pd.DataFrame,\n",
    "                 Y_df: pd.DataFrame,\n",
    "                 X_df: pd.DataFrame, \n",
    "                 mask_df: pd.DataFrame) -> Tuple[List[np.ndarray], \n",
    "                                                 List[np.ndarray],\n",
    "                                                 List[np.ndarray], \n",
    "                                                 List[str],\n",
    "                                                 List[str]]:\n",
    "    \"\"\"Transforms input dataframes to lists.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    S_df: pd.DataFrame\n",
    "        Static exogenous variables with columns ['unique_id', 'ds'] \n",
    "        and static variables.    \n",
    "    Y_df: pd.DataFrame\n",
    "        Target time series with columns ['unique_id', 'ds', 'y'].\n",
    "    X_df: pd.DataFrame\n",
    "        Exogenous time series with columns ['unique_id', 'ds', 'y'].\n",
    "    mask_df: pd.DataFrame\n",
    "        Outsample mask with columns ['unique_id', 'ds', 'sample_mask']\n",
    "        and optionally 'available_mask'.\n",
    "        Default None: constructs default mask based on ds_in_test.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Tuple of five lists:\n",
    "        - List of time series. Each element of the list is a \n",
    "          numpy array of shape (length of the time series, n_channels),\n",
    "          where n_channels = t_cols + masks.\n",
    "        - List of static variables. Each element of the list is a \n",
    "          numpy array of shape (1, n_s).\n",
    "          where n_channels = t_cols + masks.\n",
    "        - List of meta data. Each element of the list is a \n",
    "          numpy array of shape (lenght of the time series, 2) \n",
    "          and corresponds to unique_id, ds.\n",
    "        - List of temporal variables (including target and masks). \n",
    "        - List of statitc variables.\n",
    "    \"\"\"\n",
    "    # None protections\n",
    "    if X_df is None:\n",
    "        X_df = Y_df[['unique_id', 'ds']]\n",
    "    \n",
    "    if S_df is None:\n",
    "        S_df = Y_df[['unique_id']].drop_duplicates()\n",
    "    \n",
    "    # Protect order of data\n",
    "    Y = Y_df.sort_values(by=['unique_id', 'ds'], ignore_index=True).copy()\n",
    "    X = X_df.sort_values(by=['unique_id', 'ds'], ignore_index=True).copy()\n",
    "    M = mask_df.sort_values(by=['unique_id', 'ds'], ignore_index=True).copy()\n",
    "    \n",
    "    assert np.array_equal(X.unique_id.values, Y.unique_id.values), f'Mismatch in X, Y unique_ids'\n",
    "    assert np.array_equal(X.ds.values, Y.ds.values), f'Mismatch in X, Y ds'\n",
    "    assert np.array_equal(M.unique_id.values, Y.unique_id.values), f'Mismatch in M, Y unique_ids'\n",
    "    assert np.array_equal(M.ds.values, Y.ds.values), f'Mismatch in M, Y ds'\n",
    "    \n",
    "    # Create bigger grouped by dataframe G to parse\n",
    "    M = M[['available_mask', 'sample_mask']]\n",
    "    X.drop(['unique_id', 'ds'], 1, inplace=True)\n",
    "    G = Y.join(X).join(M)\n",
    "    \n",
    "    S = S_df.sort_values('unique_id')\n",
    "    \n",
    "    # time columns and static columns for future indexing\n",
    "    t_cols = list(G.columns[2:]) # avoid unique_id and ds\n",
    "    s_cols = list(S.columns[1:]) # avoid unique_id\n",
    "    \n",
    "    grouped = G.groupby('unique_id')\n",
    "    meta = G[['unique_id', 'ds']].values\n",
    "    data = G.drop(columns=['unique_id', 'ds']).values\n",
    "    sizes = grouped.size()\n",
    "    idxs = np.append(0, sizes.cumsum())\n",
    "    ts_data = []\n",
    "    meta_data = []\n",
    "    for start, end in zip(idxs[:-1], idxs[1:]):\n",
    "        ts_data.append(data[start:end])\n",
    "        meta_data.append(meta[start:end])\n",
    "        \n",
    "    if S['unique_id'].value_counts().max() > 1:\n",
    "        raise ValueError('Found duplicated unique_ids in S_df')\n",
    "    s_data = S.drop(columns='unique_id').values\n",
    "    \n",
    "    del S, Y, X, M, G\n",
    "    gc.collect()\n",
    "    \n",
    "    return ts_data, s_data, meta_data, t_cols, s_cols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@patch\n",
    "def _create_tensor(self: BaseDataset) -> Tuple[np.array, t.Tensor]:\n",
    "    \"\"\"Transforms outputs from self._df_to_lists to numpy arrays.\"\"\"\n",
    "    ts_tensor = np.zeros((self.n_series, self.n_channels, self.max_len))\n",
    "\n",
    "    len_series = np.empty(self.n_series, dtype=np.int32)\n",
    "    for idx, ts_idx in enumerate(self.ts_data):\n",
    "        # Left padded time series tensor\n",
    "        ts_tensor[idx, :, -ts_idx.shape[0]:] = ts_idx.T\n",
    "        len_series[idx] = ts_idx.shape[0]\n",
    "    \n",
    "    ts_tensor = t.Tensor(ts_tensor)\n",
    "\n",
    "    return len_series, ts_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def _get_f_idxs(self: BaseDataset, \n",
    "               cols: List[str]) -> List:\n",
    "    \"\"\"Gets indexes of exogenous variables.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    cols: List[str]\n",
    "        Interest exogenous variables.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Indexes of cols variables.\n",
    "    \"\"\"\n",
    "    # Check if cols are available f_cols and return the idxs\n",
    "    if not all(col in self.f_cols for col in cols):\n",
    "        str_cols = ', '.join(cols)\n",
    "        raise Exception(f'Some variables in {str_cols} are not available in f_cols.')\n",
    "    \n",
    "    f_idxs = [self.t_cols.index(col) for col in cols]\n",
    "\n",
    "    return f_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __getitem__(self: BaseDataset, \n",
    "                idx: Union[slice, int]) -> Dict[str, t.Tensor]:\n",
    "    \"\"\"Creates batch based on index.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    index: np.ndarray\n",
    "        Indexes of time series to consider.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Dictionary with keys:\n",
    "        - S\n",
    "        - Y\n",
    "        - X\n",
    "        - available_mask\n",
    "        - sample_mask\n",
    "        - idxs\n",
    "    \"\"\"\n",
    "    # Checks for idx\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __len__(self: BaseDataset):\n",
    "    return self.n_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def get_n_variables(self: BaseDataset) -> Tuple[int, int]:\n",
    "    \"\"\"Gets number of exogenous and static variables.\"\"\"\n",
    "    return self.n_x, self.n_s\n",
    "\n",
    "@patch\n",
    "def get_n_series(self: BaseDataset) -> int:\n",
    "    \"\"\"Gets number of time series.\"\"\"\n",
    "    return self.n_series\n",
    "\n",
    "@patch\n",
    "def get_max_len(self: BaseDataset) -> int:\n",
    "    \"\"\"Gets max len of time series.\"\"\"\n",
    "    return self.max_len\n",
    "\n",
    "@patch\n",
    "def get_n_channels(self: BaseDataset) -> int:\n",
    "    \"\"\"Gets number of channels considered.\"\"\"\n",
    "    return self.n_channels\n",
    "\n",
    "@patch\n",
    "def get_frequency(self: BaseDataset) -> str:\n",
    "    \"\"\"Gets infered frequency.\"\"\"\n",
    "    return self.frequency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_default_mask_df(Y_df: pd.DataFrame, \n",
    "                        ds_in_test: int, \n",
    "                        is_test: bool) -> pd.DataFrame:\n",
    "    \"\"\"Constructs default mask df.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    Y_df: pd.DataFrame\n",
    "        Target time series with columns ['unique_id', 'ds', 'y'].\n",
    "    ds_in_test: int\n",
    "        Numer of datestamps to use as outsample.\n",
    "    is_test: bool\n",
    "        Wheter target time series belongs to test set.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Mask DataFrame with columns \n",
    "    ['unique_id', 'ds', 'available_mask', 'sample_mask'].\n",
    "    \"\"\"\n",
    "    mask_df = Y_df[['unique_id', 'ds']].copy()\n",
    "    mask_df['available_mask'] = 1\n",
    "    mask_df['sample_mask'] = 1\n",
    "    mask_df = mask_df.set_index(['unique_id', 'ds'])\n",
    "    \n",
    "    mask_df_s = mask_df.sort_values(by=['unique_id', 'ds'])\n",
    "    zero_idx = mask_df_s.groupby('unique_id').tail(ds_in_test).index\n",
    "    mask_df.loc[zero_idx, 'sample_mask'] = 0\n",
    "    mask_df = mask_df.reset_index()\n",
    "    mask_df.index = Y_df.index\n",
    "\n",
    "    assert len(mask_df)==len(Y_df), \\\n",
    "        f'The mask_df length {len(mask_df)} is not equal to Y_df length {len(Y_df)}'\n",
    "\n",
    "    if is_test:\n",
    "        mask_df['sample_mask'] = 1 - mask_df['sample_mask']\n",
    "\n",
    "    return mask_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TimeSeries Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TimeSeriesDataset(BaseDataset):\n",
    "    \"\"\"\n",
    "    A class used to store Time Series data.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,\n",
    "                 Y_df: pd.DataFrame,\n",
    "                 input_size: int,\n",
    "                 output_size: int,\n",
    "                 X_df: Optional[pd.DataFrame] = None,\n",
    "                 S_df: Optional[pd.DataFrame] = None,\n",
    "                 f_cols: Optional[List] = None,\n",
    "                 mask_df: Optional[pd.DataFrame] = None,\n",
    "                 ds_in_test: int = 0,\n",
    "                 is_test: bool = False, \n",
    "                 complete_windows: bool = True,\n",
    "                 verbose: bool = False) -> 'TimeSeriesDataset':\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        Y_df: pd.DataFrame\n",
    "            Target time series with columns ['unique_id', 'ds', 'y'].\n",
    "        input_size: int\n",
    "            Size of the training sets.\n",
    "        output_size: int\n",
    "            Forecast horizon.\n",
    "        X_df: pd.DataFrame\n",
    "            Exogenous time series with columns ['unique_id', 'ds', 'y'].\n",
    "        S_df: pd.DataFrame\n",
    "            Static exogenous variables with columns ['unique_id', 'ds'] \n",
    "            and static variables.\n",
    "        f_cols: list\n",
    "            List of exogenous variables of the future.\n",
    "        mask_df: pd.DataFrame\n",
    "            Outsample mask with columns ['unique_id', 'ds', 'sample_mask']\n",
    "            and optionally 'available_mask'.\n",
    "            Default None: constructs default mask based on ds_in_test.\n",
    "        ds_in_test: int\n",
    "            Only used when mask_df = None.\n",
    "            Numer of datestamps to use as outsample.\n",
    "        is_test: bool\n",
    "            Only used when mask_df = None.\n",
    "            Wheter target time series belongs to test set.\n",
    "        verbose: bool\n",
    "            Wheter or not log outputs.\n",
    "        \"\"\"        \n",
    "        super(TimeSeriesDataset, self).__init__(Y_df=Y_df, input_size=input_size,\n",
    "                                                output_size=output_size,\n",
    "                                                X_df=X_df, S_df=S_df, f_cols=f_cols,\n",
    "                                                mask_df=mask_df, ds_in_test=ds_in_test,\n",
    "                                                is_test=is_test, complete_windows=complete_windows,\n",
    "                                                verbose=verbose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __getitem__(self: TimeSeriesDataset, \n",
    "                idx: Union[slice, int]) -> Dict[str, t.Tensor]:\n",
    "    \"\"\"Creates batch based on index.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    index: np.ndarray\n",
    "        Indexes of time series to consider.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Dictionary with keys:\n",
    "        - S\n",
    "        - Y\n",
    "        - X\n",
    "        - available_mask\n",
    "        - sample_mask\n",
    "        - idxs\n",
    "    \"\"\"\n",
    "    # Checks for idx\n",
    "    if isinstance(idx, int):\n",
    "        idx = [idx]\n",
    "    elif isinstance(idx, slice) or isinstance(idx, list):\n",
    "        pass\n",
    "    else:\n",
    "        raise Exception('Use slices, int or list for getitem.')\n",
    "\n",
    "    # Parse windows to elements of batch\n",
    "    S = t.Tensor(self.s_matrix[idx])\n",
    "    Y = self.ts_tensor[idx, self.t_cols.index('y'), :]\n",
    "    X = self.ts_tensor[idx, (self.t_cols.index('y') + 1):self.t_cols.index('available_mask'), :]\n",
    "    \n",
    "    available_mask = self.ts_tensor[idx, self.t_cols.index('available_mask'), :]\n",
    "    sample_mask = self.ts_tensor[idx, self.t_cols.index('sample_mask'), :]\n",
    "    ts_idxs = t.as_tensor(idx, dtype=t.long)\n",
    "\n",
    "    batch = {'S': S, 'Y': Y, 'X': X,\n",
    "             'available_mask': available_mask,\n",
    "             'sample_mask': sample_mask,\n",
    "             'idxs': ts_idxs}\n",
    "    \n",
    "    return batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Windows Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class WindowsDataset(BaseDataset):\n",
    "    \"\"\"\n",
    "    A class used to store Time Series data.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self,\n",
    "                 Y_df: pd.DataFrame,\n",
    "                 input_size: int,\n",
    "                 output_size: int,\n",
    "                 X_df: Optional[pd.DataFrame] = None,\n",
    "                 S_df: Optional[pd.DataFrame] = None,\n",
    "                 f_cols: Optional[List] = None,\n",
    "                 mask_df: Optional[pd.DataFrame] = None,\n",
    "                 ds_in_test: int = 0,\n",
    "                 is_test: bool = False,\n",
    "                 sample_freq: int = 1,\n",
    "                 complete_windows: bool = True,\n",
    "                 last_window: bool = False,\n",
    "                 verbose: bool = False) -> 'TimeSeriesDataset':\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        Y_df: pd.DataFrame\n",
    "            Target time series with columns ['unique_id', 'ds', 'y'].\n",
    "        input_size: int\n",
    "            Size of the training sets.\n",
    "        output_size: int\n",
    "            Forecast horizon.\n",
    "        X_df: pd.DataFrame\n",
    "            Exogenous time series with columns ['unique_id', 'ds', 'y'].\n",
    "        S_df: pd.DataFrame\n",
    "            Static exogenous variables with columns ['unique_id', 'ds'] \n",
    "            and static variables.\n",
    "        f_cols: list\n",
    "            List of exogenous variables of the future.\n",
    "        mask_df: pd.DataFrame\n",
    "            Outsample mask with columns ['unique_id', 'ds', 'sample_mask']\n",
    "            and optionally 'available_mask'.\n",
    "            Default None: constructs default mask based on ds_in_test.\n",
    "        ds_in_test: int\n",
    "            Only used when mask_df = None.\n",
    "            Numer of datestamps to use as outsample.\n",
    "        is_test: bool\n",
    "            Only used when mask_df = None.\n",
    "            Wheter target time series belongs to test set.\n",
    "        last_window: bool\n",
    "            Only used for forecast (test)\n",
    "            Wheter the dataset will include only last window for each time serie.\n",
    "        verbose: bool\n",
    "            Wheter or not log outputs.\n",
    "        \"\"\"        \n",
    "        super(WindowsDataset, self).__init__(Y_df=Y_df, input_size=input_size,\n",
    "                                             output_size=output_size,\n",
    "                                             X_df=X_df, S_df=S_df, f_cols=f_cols,\n",
    "                                             mask_df=mask_df, ds_in_test=ds_in_test,\n",
    "                                             is_test=is_test, complete_windows=complete_windows,\n",
    "                                             verbose=verbose)\n",
    "        # WindowsDataset parameters\n",
    "        self.windows_size = self.input_size + self.output_size\n",
    "        self.padding = (self.input_size, self.output_size)\n",
    "        self.sample_freq = sample_freq\n",
    "        self.last_window = last_window"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def _create_windows_tensor(self: WindowsDataset, \n",
    "                           idx: slice) -> Tuple[t.Tensor, t.Tensor, t.Tensor]:\n",
    "    \"\"\"Creates windows of size windows_size from\n",
    "    the ts_tensor of the TimeSeriesDataset filtered by\n",
    "    window_sampling_limit and ts_idxs. The step of each window\n",
    "    is defined by idx_to_sample_freq.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    index: slice\n",
    "        Indexes of time series to consider.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Tuple of three elements:\n",
    "        - Windows tensor of shape (windows, channels, input_size + output_size)\n",
    "        - Static variables tensor of shape (windows * series, n_static)\n",
    "        - Time Series indexes for each window.\n",
    "    \"\"\"\n",
    "    # Default ts_idxs=ts_idxs sends all the data, otherwise filters series   \n",
    "    tensor = self.ts_tensor[idx, :, self.first_ds:]\n",
    "\n",
    "    padder = t.nn.ConstantPad1d(padding=self.padding, value=0)\n",
    "    tensor = padder(tensor)\n",
    "\n",
    "    # Creating rolling windows and 'flattens' them\n",
    "    windows = tensor.unfold(dimension=-1, \n",
    "                            size=self.windows_size, \n",
    "                            step=self.sample_freq)\n",
    "    # n_serie, n_channel, n_time, window_size -> n_serie, n_time, n_channel, window_size\n",
    "    windows = windows.permute(0, 2, 1, 3)\n",
    "    windows = windows.reshape(-1, self.n_channels, self.windows_size)\n",
    "    \n",
    "    # Broadcast s_matrix: This works because unfold in windows_tensor, orders: serie, time\n",
    "    \n",
    "    ts_idxs = self.ts_idxs[idx]\n",
    "    n_ts = len(ts_idxs)\n",
    "    windows_per_serie = len(windows) / n_ts\n",
    "    \n",
    "    ts_idxs = ts_idxs.repeat(repeats=windows_per_serie)\n",
    "    s_matrix = self.s_matrix[idx]\n",
    "    s_matrix = s_matrix.repeat(repeats=windows_per_serie, axis=0)\n",
    "    \n",
    "    s_matrix = t.Tensor(s_matrix)\n",
    "    ts_idxs = t.as_tensor(ts_idxs, dtype=t.long)\n",
    "\n",
    "    windows_idxs = self._get_sampleable_windows_idxs(ts_windows_flatten=windows,\n",
    "                                                     ts_idxs=ts_idxs)\n",
    "\n",
    "    # Raise error if nothing to sample from\n",
    "    if not windows_idxs.size:\n",
    "        raise Exception(\n",
    "            f'Time Series {idx} are not sampleable. '\n",
    "            'Check the data, masks, window_sampling_limit, '\n",
    "            'input_size, output_size, masks.'\n",
    "        )\n",
    "\n",
    "    # Index the windows and s_matrix tensors of batch\n",
    "    windows = windows[windows_idxs]\n",
    "    s_matrix = s_matrix[windows_idxs]\n",
    "    ts_idxs = ts_idxs[windows_idxs]\n",
    "\n",
    "    return windows, s_matrix, ts_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "#TODO: do we want complete? inputs seems irrelevant, NBEATS dont use it, for now is our only model\n",
    "def _get_sampleable_windows_idxs(self: WindowsDataset, \n",
    "                                 ts_windows_flatten: t.Tensor,\n",
    "                                 ts_idxs: t.Tensor) -> np.ndarray:\n",
    "    \"\"\"Gets indexes of windows that fulfills conditions.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    ts_windows_flatten: t.Tensor\n",
    "        Tensor of shape (windows, n_channels, windows_size)\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Numpy array of indexes of ts_windows_flatten that \n",
    "    fulfills conditions.\n",
    "    \n",
    "    Notes\n",
    "    -----\n",
    "    \"\"\"\n",
    "\n",
    "    if self.last_window:\n",
    "        _, idxs_counts = t.unique(ts_idxs, return_counts=True)\n",
    "        last_idxs = idxs_counts.cumsum(0) - 1\n",
    "        last_idxs = last_idxs.numpy()\n",
    "\n",
    "        return last_idxs\n",
    "\n",
    "    else:\n",
    "        sample_condition = ts_windows_flatten[:, self.t_cols.index('sample_mask'), -self.output_size:]\n",
    "        sample_condition = t.sum(sample_condition, axis=1)\n",
    "        sample_condition = (sample_condition > 0) * 1           \n",
    "        sampling_idx = t.nonzero(sample_condition > 0)\n",
    "        sampling_idx = sampling_idx.flatten().numpy()\n",
    "        \n",
    "        return sampling_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __getitem__(self: WindowsDataset, \n",
    "                idx: Union[slice, int]) -> Dict[str, t.Tensor]:\n",
    "    \"\"\"Creates batch based on index.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    index: np.ndarray\n",
    "        Indexes of time series to consider.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    Dictionary with keys:\n",
    "        - S\n",
    "        - Y\n",
    "        - X\n",
    "        - available_mask\n",
    "        - sample_mask\n",
    "        - idxs\n",
    "    \"\"\"\n",
    "    # Checks for idx\n",
    "    if isinstance(idx, int):\n",
    "        idx = [idx]\n",
    "    elif isinstance(idx, slice) or isinstance(idx, list):\n",
    "        pass\n",
    "    else:\n",
    "        raise Exception('Use slices, int or list for getitem.')\n",
    "\n",
    "    # Create windows for each sampled ts and sample random unmasked windows from each ts\n",
    "    windows, S, ts_idxs = self._create_windows_tensor(idx=idx)\n",
    "\n",
    "    # Parse windows to elements of batch\n",
    "    Y = windows[:, self.t_cols.index('y'), :]\n",
    "    X = windows[:, (self.t_cols.index('y') + 1):self.t_cols.index('available_mask'), :]\n",
    "    available_mask = windows[:, self.t_cols.index('available_mask'), :]\n",
    "    sample_mask = windows[:, self.t_cols.index('sample_mask'), :]\n",
    "\n",
    "    batch = {'S': S, 'Y': Y, 'X': X,\n",
    "             'available_mask': available_mask,\n",
    "             'sample_mask': sample_mask,\n",
    "             'idxs': ts_idxs}\n",
    "    \n",
    "    return batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Default mask example and tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_default_mask(Y_df, ds_in_test, is_test):\n",
    "    mask_df = get_default_mask_df(Y_df, ds_in_test, is_test)\n",
    "    assert Y_df.index.equals(mask_df.index), 'Unmatching index bewteen Y_df and mask_df'\n",
    "    \n",
    "    for uid, df in mask_df.groupby('unique_id'):\n",
    "        len_ts = df.shape[0]\n",
    "        expected_sample_mask = np.ones(len_ts)\n",
    "        expected_sample_mask[-ds_in_test:] = 0\n",
    "        if is_test: \n",
    "            expected_sample_mask = 1 - expected_sample_mask\n",
    "        expected_available_mask = np.ones(len_ts)\n",
    "        \n",
    "        sample_mask = df['sample_mask'].values\n",
    "        available_mask = df['available_mask'].values\n",
    "        \n",
    "        assert np.array_equal(sample_mask, expected_sample_mask), (\n",
    "            f'Error for sample mask for time series {uid}'\n",
    "        )\n",
    "        \n",
    "        assert np.array_equal(available_mask, expected_available_mask), (\n",
    "            f'Error for available mask for time series {uid}'\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test for synthtetic time series data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nixtlats.data.utils import create_synthetic_tsdata\n",
    "\n",
    "Y_df, X_df, S_df = create_synthetic_tsdata()\n",
    "ds_in_test = 2\n",
    "is_test = False\n",
    "test_default_mask(Y_df, ds_in_test, is_test)\n",
    "test_default_mask(Y_df, ds_in_test, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example and test for datasets with two time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fb7fd6d1ad0>]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 6236 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from nixtlats.data.datasets.epf import EPF, EPFInfo\n",
    "\n",
    "Y_df, X_df, S_df = EPF.load_groups(directory='data', groups=['NP', 'PJM'])\n",
    "test_default_mask(Y_df, ds_in_test=728 * 24, is_test=False)\n",
    "mask_df = get_default_mask_df(Y_df=Y_df, ds_in_test=728 * 24, is_test=False)\n",
    "\n",
    "plt.plot(mask_df.sample_mask.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7fb7fbc96990>]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 6236 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "mask_df = get_default_mask_df(Y_df=Y_df, ds_in_test=728 * 24, is_test=True)\n",
    "\n",
    "plt.plot(mask_df.sample_mask.values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test for datasets with more than two time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nixtlats.data.datasets.tourism import Tourism, TourismInfo\n",
    "\n",
    "meta = TourismInfo['Yearly']\n",
    "Y_df, *_ = Tourism.load(directory='data', group=meta.name)\n",
    "test_default_mask(Y_df, ds_in_test=meta.horizon, is_test=False)\n",
    "test_default_mask(Y_df, ds_in_test=meta.horizon, is_test=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def instantiate_datasets(Y_df, S_df, X_df, f_cols=None, \n",
    "                         ds_in_test=0, is_test=False,\n",
    "                         input_size=15,\n",
    "                         output_size=1,\n",
    "                         complete_windows=False,\n",
    "                         sample_freq=1):\n",
    "    mask_df = get_default_mask_df(Y_df=Y_df, ds_in_test=ds_in_test, is_test=is_test)\n",
    "                        \n",
    "    ts_dataset = TimeSeriesDataset(Y_df=Y_df, S_df=S_df, X_df=X_df, f_cols=f_cols, \n",
    "                                   mask_df=mask_df,\n",
    "                                   input_size=input_size,\n",
    "                                   output_size=output_size,\n",
    "                                   complete_windows=complete_windows)\n",
    "    \n",
    "    wd_dataset = WindowsDataset(Y_df=Y_df, S_df=S_df, X_df=X_df, f_cols=f_cols, \n",
    "                                   mask_df=mask_df,\n",
    "                                   input_size=input_size,\n",
    "                                   output_size=output_size,\n",
    "                                   sample_freq=sample_freq,\n",
    "                                   complete_windows=complete_windows)\n",
    "    \n",
    "    return ts_dataset, wd_dataset, mask_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_dataset_attrs(Y_df, S_df, X_df, f_cols, ds_in_test, is_test):\n",
    "    # This set catches mistmaches between Y_df and ts_tensor\n",
    "    ts_dataset, wd_dataset, mask_df = instantiate_datasets(Y_df=Y_df, S_df=S_df, X_df=X_df, \n",
    "                                           f_cols=f_cols, ds_in_test=ds_in_test, \n",
    "                                           is_test=is_test)\n",
    "    \n",
    "    dfs = [Y_df, X_df, mask_df]\n",
    "    dfs = [df.set_index(['unique_id', 'ds']) for df in dfs]\n",
    "    dfs = dfs[0].join(dfs[1:])\n",
    "    \n",
    "    #Temporal variables\n",
    "    for dataset in [ts_dataset, wd_dataset]:\n",
    "        for idx_ts, (uid, df) in enumerate(dfs.groupby('unique_id')):\n",
    "            len_ts = dataset.len_series[idx_ts]\n",
    "\n",
    "            for col in dataset.t_cols:\n",
    "                ts = t.Tensor(df[col].values)\n",
    "                idx_tensor = dataset.t_cols.index(col)\n",
    "                ts_tensor = dataset.ts_tensor[idx_ts, idx_tensor, -len_ts:]\n",
    "\n",
    "                assert np.array_equal(ts, ts_tensor), (\n",
    "                    f'Error with time series {uid} and col {col} (idx={idx_ts}).'\n",
    "                )\n",
    "\n",
    "        #Static variables\n",
    "        for idx_ts, (uid, df) in enumerate(S_df.groupby('unique_id')):\n",
    "            len_ts = dataset.len_series[idx_ts]\n",
    "\n",
    "            s = df[dataset.s_cols].values\n",
    "            s_matrix = dataset.s_matrix[[idx_ts]]\n",
    "\n",
    "            assert np.array_equal(s, s_matrix), (\n",
    "                f'Error with static variables for time series {uid} (idx={idx_ts})'\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_get_f_idxs(Y_df, S_df, X_df, f_cols, ds_in_test, is_test, expected_f_idxs):\n",
    "    ts_dataset, wd_dataset, mask_df = instantiate_datasets(Y_df=Y_df, S_df=S_df, X_df=X_df, \n",
    "                                           f_cols=f_cols, ds_in_test=ds_in_test, \n",
    "                                           is_test=is_test)\n",
    "    \n",
    "    assert ts_dataset._get_f_idxs(f_cols) == expected_f_idxs\n",
    "    assert wd_dataset._get_f_idxs(f_cols) == expected_f_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_ts_tensor(Y_df, S_df, X_df, f_cols, ds_in_test, is_test, \n",
    "                   input_size, output_size, ts_idxs):\n",
    "    ts_dataset, wd_dataset, mask_df = instantiate_datasets(Y_df=Y_df, S_df=S_df, X_df=X_df, \n",
    "                                           f_cols=f_cols, ds_in_test=ds_in_test, \n",
    "                                           is_test=is_test,\n",
    "                                           input_size=input_size,\n",
    "                                           output_size=output_size)\n",
    "    \n",
    "    for dataset in [ts_dataset, wd_dataset]:\n",
    "        min_len = min(dataset.len_series)\n",
    "        dfs = [Y_df, X_df, mask_df]\n",
    "        dfs = [df.set_index(['unique_id', 'ds']) for df in dfs]\n",
    "        dfs = dfs[0].join(dfs[1:])\n",
    "        \n",
    "        # This process only works for balanced datasets.\n",
    "        \n",
    "        n_ts = Y_df['unique_id'].unique().shape[0]\n",
    "        n_x = dfs.columns.shape[0]\n",
    "        idxs = range(n_ts) if ts_idxs is None else ts_idxs\n",
    "\n",
    "        e_filtered_tensor = t.Tensor(dfs.values.reshape((n_ts, min_len, n_x))[idxs])\n",
    "        e_filtered_tensor = np.swapaxes(e_filtered_tensor, 2, 1)\n",
    "        filtered_tensor = dataset.ts_tensor[ts_idxs, :, dataset.first_ds:]\n",
    "\n",
    "        assert np.array_equal(e_filtered_tensor, filtered_tensor), (\n",
    "            \"Expected and dataset filtered_tensor are different. Check.\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy.lib.stride_tricks import sliding_window_view\n",
    "\n",
    "# This test only works for the synthetic dataset constructed \n",
    "# using create_synthetic_tsdata\n",
    "# and for the 21 time series\n",
    "def test_batch_construction_windows(Y_df, S_df, X_df, ds_in_test, \n",
    "                                    is_test, input_size, output_size, sample_freq, \n",
    "                                    ):\n",
    "    \"\"\"Test to verify that the batch (of windows) is well constructed.\"\"\"\n",
    "\n",
    "    _, dataset, mask_df = instantiate_datasets(Y_df=Y_df, S_df=S_df, X_df=X_df, \n",
    "                                               ds_in_test=ds_in_test, is_test=is_test,\n",
    "                                               input_size=input_size,\n",
    "                                               output_size=output_size,\n",
    "                                               sample_freq=sample_freq)\n",
    "    \n",
    "    windows_size = input_size + output_size\n",
    "    max_len = Y_df.groupby('unique_id').size().max()\n",
    "    \n",
    "    windows = dataset[20]['Y'].numpy()\n",
    "    \n",
    "    #Expected windows\n",
    "    uid = Y_df['unique_id'].unique()[20]\n",
    "    Y_original = Y_df.query('unique_id == @uid')['y'].values\n",
    "    size = Y_original.size\n",
    "    Y_sample = np.zeros(max_len)\n",
    "    Y_sample[-size:] = Y_original\n",
    "    \n",
    "    Y_sample = np.pad(Y_sample, (input_size+100, output_size))\n",
    "\n",
    "    e_windows = sliding_window_view(Y_sample, window_shape=windows_size)\n",
    "    e_windows = e_windows[0:-(ds_in_test+output_size-1):sample_freq] #-1 for at least one available sample\n",
    "\n",
    "    # This test works assuming there are no series with all values of zero.\n",
    "    sampleable_windows_idxs = np.where((e_windows > 0).sum(1) >= 1)[0]     \n",
    "    e_windows = e_windows[sampleable_windows_idxs]\n",
    "    \n",
    "    #Comparison\n",
    "    assert np.array_equal(windows, e_windows), (\n",
    "        'Expected and actual windows are different'\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy.lib.stride_tricks import sliding_window_view\n",
    "\n",
    "# This test only works for the synthetic dataset constructed \n",
    "# using create_synthetic_tsdata\n",
    "# and for the 21 time series\n",
    "def test_batch_construction_ts(Y_df, S_df, X_df, ds_in_test, \n",
    "                                    is_test, input_size, output_size, sample_freq, \n",
    "                                    ):\n",
    "    \"\"\"Test to verify that the batch (of windows) is well constructed.\"\"\"\n",
    "\n",
    "    dataset, _, mask_df = instantiate_datasets(Y_df=Y_df, S_df=S_df, X_df=X_df, \n",
    "                                               ds_in_test=ds_in_test, is_test=is_test,\n",
    "                                               input_size=input_size,\n",
    "                                               output_size=output_size,\n",
    "                                               sample_freq=sample_freq)\n",
    "    \n",
    "    batch = dataset[20]['Y'].numpy()\n",
    "    max_len = Y_df.groupby('unique_id').size().max()\n",
    "    \n",
    "    #Expected windows\n",
    "    uid = Y_df['unique_id'].unique()[20]\n",
    "    Y_original = Y_df.query('unique_id == @uid')['y'].values\n",
    "    size = Y_original.size\n",
    "    e_batch = np.zeros(max_len)\n",
    "    e_batch[-size:] = Y_original\n",
    "    e_batch = e_batch[None,:]\n",
    "\n",
    "    #Comparison\n",
    "    assert np.array_equal(batch, e_batch), (\n",
    "        'Expected and actual windows are different'\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test for not sorted datasets with more than two time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nixtlats.data.utils import create_synthetic_tsdata\n",
    "\n",
    "Y_df, X_df, S_df = create_synthetic_tsdata()\n",
    "ds_in_test = 2\n",
    "is_test = False\n",
    "f_cols = ['future_1']\n",
    "expected_f_idxs = [2]\n",
    "len_sample_chunks = 15 #only for ESRNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset_attrs(Y_df, S_df, X_df, f_cols=f_cols, ds_in_test=ds_in_test, is_test=is_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_get_f_idxs(Y_df, S_df, X_df, f_cols=f_cols, ds_in_test=ds_in_test, is_test=is_test, \n",
    "                expected_f_idxs=expected_f_idxs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Expected error for non-sorted datasets \n",
    "\n",
    "For the `ts_tensor` attribute from the `dataset` and `Y_df` to have the same order it is necessary that `Y_df` is sorted by `unique_id` and `ds`. This test (expected error) proves that for unordered data the order is different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _fail_non_sorted(): \n",
    "    test_ts_tensor(Y_df, S_df, X_df, f_cols=f_cols, \n",
    "                   ds_in_test=ds_in_test, \n",
    "                   is_test=is_test,\n",
    "                   ts_idxs=[1, 0], \n",
    "                   output_size=ds_in_test)\n",
    "test_fail(_fail_non_sorted)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Test using sorted synthtic ts data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nixtlats.data.utils import create_synthetic_tsdata\n",
    "\n",
    "Y_df, X_df, S_df = create_synthetic_tsdata(sort=True)\n",
    "ds_in_test = 2\n",
    "is_test = False\n",
    "f_cols = ['future_1']\n",
    "expected_f_idxs = [2]\n",
    "len_sample_chunks = 15 #only for ESRNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_batch_construction_windows(Y_df, S_df, X_df, ds_in_test=ds_in_test,\n",
    "                                is_test=is_test, input_size=5, output_size=2,\n",
    "                                sample_freq=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_batch_construction_ts(Y_df, S_df, X_df, ds_in_test=ds_in_test,\n",
    "                           is_test=is_test, input_size=5, output_size=2,\n",
    "                           sample_freq=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test for already sorted datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nixtlats.data.datasets.epf import EPF, EPFInfo\n",
    "\n",
    "Y_df, X_df, S_df = EPF.load_groups(directory='data', groups=['NP', 'PJM'])\n",
    "f_cols = ['Exogenous1', 'Exogenous2']\n",
    "ds_in_test = 728 * 24\n",
    "is_test = True\n",
    "expected_f_idxs = [1, 2] #after y column"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset_attrs(Y_df, S_df, X_df, f_cols=f_cols, ds_in_test=ds_in_test, is_test=is_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ts_tensor(Y_df, S_df, X_df, f_cols=f_cols, ds_in_test=ds_in_test, is_test=is_test,\n",
    "               input_size=ds_in_test, output_size=ds_in_test, ts_idxs=[1, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_get_f_idxs(Y_df, S_df, X_df, f_cols=f_cols, ds_in_test=ds_in_test, is_test=is_test, \n",
    "                expected_f_idxs=expected_f_idxs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test for already sorted datasets with more than two time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nixtlats.data.datasets.tourism import Tourism, TourismInfo\n",
    "\n",
    "meta = TourismInfo['Yearly']\n",
    "df, *_ = Tourism.load(directory='./data', group=meta.name)\n",
    "df['day_of_week'] = df['ds'].dt.day_of_week\n",
    "df['id_ts'] = df['unique_id'].astype('category').cat.codes\n",
    "\n",
    "Y_df = df.filter(items=['unique_id', 'ds', 'y'])\n",
    "X_df = df.filter(items=['unique_id', 'ds', 'day_of_week'])\n",
    "S_df = df.filter(items=['unique_id', 'id_ts']).drop_duplicates().reset_index(drop=True)\n",
    "Y_df = Y_df.groupby('unique_id').tail(11)\n",
    "X_df = X_df.groupby('unique_id').tail(11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset_attrs(Y_df, S_df, X_df, f_cols=[], ds_in_test=meta.horizon, is_test=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ts_tensor(Y_df, S_df, X_df, f_cols=[], ds_in_test=meta.horizon, is_test=False,\n",
    "               input_size=5, \n",
    "               output_size=meta.horizon, ts_idxs=[1, 7, 10, 15])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_get_f_idxs(Y_df, S_df, X_df, f_cols=[], ds_in_test=ds_in_test, is_test=is_test, \n",
    "                expected_f_idxs=[])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nixtla",
   "language": "python",
   "name": "nixtla"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
