{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_exp data.datasets.utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset Utils\n",
    "\n",
    "> Utils for dataset processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import logging\n",
    "import requests\n",
    "import subprocess\n",
    "import zipfile\n",
    "from pathlib import Path\n",
    "from dataclasses import dataclass\n",
    "from typing import Dict, List, Optional, Tuple, Union\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def download_file(directory: str, source_url: str, decompress: bool = False) -> None:\n",
    "    \"\"\"Download data from source_ulr inside directory.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    directory: str, Path\n",
    "        Custom directory where data will be downloaded.\n",
    "    source_url: str\n",
    "        URL where data is hosted.\n",
    "    decompress: bool\n",
    "        Wheter decompress downloaded file. Default False.\n",
    "    \"\"\"\n",
    "    if isinstance(directory, str):\n",
    "        directory = Path(directory)\n",
    "    directory.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    filename = source_url.split('/')[-1]\n",
    "    filepath = Path(f'{directory}/{filename}')\n",
    "\n",
    "    # Streaming, so we can iterate over the response.\n",
    "    headers = {'User-Agent': 'Mozilla/5.0'}\n",
    "    r = requests.get(source_url, stream=True, headers=headers)\n",
    "    # Total size in bytes.\n",
    "    total_size = int(r.headers.get('content-length', 0))\n",
    "    block_size = 1024 #1 Kibibyte\n",
    "\n",
    "    t = tqdm(total=total_size, unit='iB', unit_scale=True)\n",
    "    with open(filepath, 'wb') as f:\n",
    "        for data in r.iter_content(block_size):\n",
    "            t.update(len(data))\n",
    "            f.write(data)\n",
    "            f.flush()\n",
    "    t.close()\n",
    "\n",
    "    if total_size != 0 and t.n != total_size:\n",
    "        logger.error('ERROR, something went wrong downloading data')\n",
    "\n",
    "    size = filepath.stat().st_size\n",
    "    logger.info(f'Successfully downloaded {filename}, {size}, bytes.')\n",
    "\n",
    "    if decompress:\n",
    "        if '.zip' in filepath.suffix:\n",
    "            logger.info('Decompressing zip file...')\n",
    "            with zipfile.ZipFile(filepath, 'r') as zip_ref:\n",
    "                zip_ref.extractall(directory)\n",
    "        else:\n",
    "            from patoolib import extract_archive\n",
    "            extract_archive(filepath, outdir=directory)\n",
    "        logger.info(f'Successfully decompressed {filepath}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@dataclass\n",
    "class Info:\n",
    "    \"\"\"\n",
    "    Info Dataclass of datasets.\n",
    "    Args:\n",
    "        groups (Tuple): Tuple of str groups\n",
    "        class_groups (Tuple): Tuple of dataclasses.\n",
    "    \"\"\"\n",
    "    groups: Tuple[str]\n",
    "    class_groups: Tuple[dataclass]\n",
    "\n",
    "    def get_group(self, group: str):\n",
    "        \"\"\"Gets dataclass of group.\"\"\"\n",
    "        if group not in self.groups:\n",
    "            raise Exception(f'Unkown group {group}')\n",
    "\n",
    "        return self.class_groups[self.groups.index(group)]\n",
    "    \n",
    "    def __getitem__(self, group: str):\n",
    "        \"\"\"Gets dataclass of group.\"\"\"\n",
    "        if group not in self.groups:\n",
    "            raise Exception(f'Unkown group {group}')\n",
    "\n",
    "        return self.class_groups[self.groups.index(group)]\n",
    "    \n",
    "    def __iter__(self):\n",
    "        for group in self.groups:\n",
    "            yield group, self.get_group(group)\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@dataclass\n",
    "class TimeSeriesDataclass:\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        S (pd.DataFrame): DataFrame of static features of shape\n",
    "            (n_time_series, n_features).\n",
    "        X (pd.DataFrame): DataFrame of exogenous variables of shape\n",
    "            (sum n_periods_i for i=1..n_time_series, n_exogenous).\n",
    "        Y (pd.DataFrame): DataFrame of target variable of shape\n",
    "            (sum n_periods_i for i=1..n_time_series, 1).\n",
    "        idx_categorical_static (list, optional): List of categorical indexes\n",
    "            of S.\n",
    "        group (str, optional): Group name if applies.\n",
    "            Example: 'Yearly'\n",
    "    \"\"\"\n",
    "    S: pd.DataFrame\n",
    "    X: pd.DataFrame\n",
    "    Y: pd.DataFrame\n",
    "    idx_categorical_static: Optional[List] = None\n",
    "    group: Union[str, List[str]] = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# US Calendar Utils\n",
    "- [US Official Federal definitions](https://www.opm.gov/policy-data-oversight/pay-leave/federal-holidays/)\n",
    "- [Pandas calendar utility](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import pandas as pd\n",
    "from pandas.tseries.holiday import (\n",
    "    AbstractHolidayCalendar,\n",
    "    Holiday,\n",
    "    USMartinLutherKingJr,\n",
    "    USPresidentsDay,\n",
    "    USMemorialDay,\n",
    "    USLaborDay,\n",
    "    USColumbusDay,\n",
    "    USThanksgivingDay,\n",
    "    nearest_workday\n",
    ")\n",
    "\n",
    "US_FEDERAL_HOLIDAYS = {'new_year': Holiday(\"New Years Day\", month=1, day=1, observance=nearest_workday),\n",
    "                       'martin_luther_king': USMartinLutherKingJr,\n",
    "                       'presidents': USPresidentsDay,\n",
    "                       'memorial': USMemorialDay,\n",
    "                       'independence': Holiday(\"July 4th\", month=7, day=4, observance=nearest_workday),\n",
    "                       'labor': USLaborDay,\n",
    "                       'columbus': USColumbusDay,\n",
    "                       'veterans': Holiday(\"Veterans Day\", month=11, day=11, observance=nearest_workday),\n",
    "                       'thanksgiving': USThanksgivingDay,\n",
    "                       'christmas': Holiday(\"Christmas\", month=12, day=25, observance=nearest_workday)}\n",
    "\n",
    "def get_holiday_dates(holiday, dates):\n",
    "    start_date = min(dates) + pd.DateOffset(days=-366)\n",
    "    end_date = max(dates) + pd.DateOffset(days=366)\n",
    "    holiday_calendar = AbstractHolidayCalendar(rules=[US_FEDERAL_HOLIDAYS[holiday]])\n",
    "    holiday_dates = holiday_calendar.holidays(start=start_date, end=end_date)\n",
    "    return np.array(holiday_dates)\n",
    "\n",
    "def holiday_kernel(holiday, dates):\n",
    "    # Get holidays around dates\n",
    "    dates = pd.DatetimeIndex(dates)\n",
    "    dates_np = np.array(dates).astype('datetime64[D]')\n",
    "    holiday_dates = get_holiday_dates(holiday, dates)\n",
    "    holiday_dates_np = np.array(pd.DatetimeIndex(holiday_dates)).astype('datetime64[D]')\n",
    "    \n",
    "    # Compute day distance to holiday\n",
    "    nearest_holiday_idx = np.expand_dims(dates_np, axis=1) - np.expand_dims(holiday_dates_np, axis=0)\n",
    "    nearest_holiday_idx = np.argmin(np.abs(nearest_holiday_idx), axis=1)\n",
    "    nearest_holiday = pd.DatetimeIndex([holiday_dates[idx] for idx in nearest_holiday_idx])\n",
    "    holiday_diff = (dates - nearest_holiday).days.values \n",
    "    return holiday_diff\n",
    "\n",
    "def create_calendar_variables(X_df: pd.DataFrame):\n",
    "    X_df['day_of_year'] = X_df.ds.dt.dayofyear\n",
    "    X_df['day_of_week'] = X_df.ds.dt.dayofweek\n",
    "    X_df['hour'] = X_df.ds.dt.hour\n",
    "    return X_df\n",
    "\n",
    "def create_us_holiday_distance_variables(X_df: pd.DataFrame):\n",
    "    dates = X_df.ds.dt.date\n",
    "    for holiday in US_FEDERAL_HOLIDAYS.keys():\n",
    "        X_df[f'holiday_dist_{holiday}'] = holiday_kernel(holiday=holiday, \n",
    "                                                         dates=dates)\n",
    "    return X_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 42424 characters]\n",
      "text/plain": [
       "<Figure size 720x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "ds = pd.date_range(start='2010-01-01', end='2012-12-31')\n",
    "holiday_dist_new_year = holiday_kernel(holiday='new_year', dates=ds)\n",
    "holiday_dist_independence = holiday_kernel(holiday='independence', dates=ds)\n",
    "\n",
    "fig = plt.figure(figsize=(10,4))\n",
    "plt.plot(ds, holiday_dist_new_year, label='new_year')\n",
    "plt.plot(ds, holiday_dist_independence, label='independence')\n",
    "plt.plot(ds, np.zeros(len(ds)))\n",
    "plt.title('Holiday Kernels')\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
