{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n",
      "/content/drive/MyDrive/nixtlats\n"
     ]
    }
   ],
   "source": [
    "# import os\n",
    "# from google.colab import drive\n",
    "# drive.mount('/content/drive')\n",
    "# os.chdir('./drive/MyDrive/nixtlats')\n",
    "# print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install pytorch-lightning\n",
    "# !pip install torchinfo\n",
    "# !pip install fastcore\n",
    "# !pip install s3fs\n",
    "# !pip install patool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.nbeats.nbeats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# N-BEATS: Neural Basis Expansion Analysis\n",
    "\n",
    "> API details.\n",
    "\n",
    "The N-BEATS model decomposes the objective signal by performing separate local nonlinear projections of the target data onto basis functions across its different blocks. Each block of the architecture consists of a multi layer perceptron that learns expansion coefficients for the backcast and forecast elements. The backcast model is used to clean the inputs of subsequent blocks, while the forecasts are summed to compose the final prediction. The blocks are grouped in stacks. Each of the potentially multiple stacks specializes in a different variant of basis functions. Depending on the basis functions the outputs of the model can be interpretable. The [original model](https://github.com/ElementAI/N-BEATS) is implemented in both pytorch and tensorflow.\n",
    "\n",
    "[Oreshkin, B. N., Carpov, D., Chapados, N., & Bengio, Y. (2020). N-BEATS: neural basis expansion analysisfor interpretable time series forecasting.  In 8th  International Conference on Learning Representations, ICLR 2020.](https://openreview.net/forum?id=r1ecqn4YwB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "import torch as t\n",
    "import torch.nn as nn\n",
    "\n",
    "from typing import Tuple\n",
    "from functools import partial\n",
    "\n",
    "from nixtlats.models.components.tcn import _TemporalConvNet\n",
    "from nixtlats.models.components.common import Chomp1d, RepeatVector\n",
    "from nixtlats.losses.utils import LossFunction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class _StaticFeaturesEncoder(nn.Module):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super(_StaticFeaturesEncoder, self).__init__()\n",
    "        layers = [nn.Dropout(p=0.5),\n",
    "                  nn.Linear(in_features=in_features, out_features=out_features),\n",
    "                  nn.ReLU()]\n",
    "        self.encoder = nn.Sequential(*layers)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        return x\n",
    "\n",
    "class _sEncoder(nn.Module):\n",
    "    def __init__(self, in_features, out_features, n_time_in):\n",
    "        super(_sEncoder, self).__init__()\n",
    "        layers = [nn.Dropout(p=0.5),\n",
    "                  nn.Linear(in_features=in_features, out_features=out_features),\n",
    "                  nn.ReLU()]\n",
    "        self.encoder = nn.Sequential(*layers)\n",
    "        self.repeat = RepeatVector(repeats=n_time_in)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # Encode and repeat values to match time\n",
    "        x = self.encoder(x)\n",
    "        x = self.repeat(x) # [N,S_out] -> [N,S_out,T]\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class IdentityBasis(nn.Module):\n",
    "    def __init__(self, backcast_size: int, forecast_size: int):\n",
    "        super().__init__()\n",
    "        self.forecast_size = forecast_size\n",
    "        self.backcast_size = backcast_size\n",
    " \n",
    "    def forward(self, theta: t.Tensor, insample_x_t: t.Tensor, outsample_x_t: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:\n",
    "        backcast = theta[:, :self.backcast_size]\n",
    "        forecast = theta[:, -self.forecast_size:]\n",
    "        return backcast, forecast\n",
    "\n",
    "class TrendBasis(nn.Module):\n",
    "    def __init__(self, degree_of_polynomial: int, backcast_size: int, forecast_size: int):\n",
    "        super().__init__()\n",
    "        polynomial_size = degree_of_polynomial + 1\n",
    "        self.backcast_basis = nn.Parameter(\n",
    "            t.tensor(np.concatenate([np.power(np.arange(backcast_size, dtype=float) / backcast_size, i)[None, :]\n",
    "                                    for i in range(polynomial_size)]), dtype=t.float32), requires_grad=False)\n",
    "        self.forecast_basis = nn.Parameter(\n",
    "            t.tensor(np.concatenate([np.power(np.arange(forecast_size, dtype=float) / forecast_size, i)[None, :]\n",
    "                                    for i in range(polynomial_size)]), dtype=t.float32), requires_grad=False)\n",
    "    \n",
    "    def forward(self, theta: t.Tensor, insample_x_t: t.Tensor, outsample_x_t: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:\n",
    "        cut_point = self.forecast_basis.shape[0]\n",
    "        backcast = t.einsum('bp,pt->bt', theta[:, cut_point:], self.backcast_basis)\n",
    "        forecast = t.einsum('bp,pt->bt', theta[:, :cut_point], self.forecast_basis)\n",
    "        return backcast, forecast\n",
    "\n",
    "class SeasonalityBasis(nn.Module):\n",
    "    def __init__(self, harmonics: int, backcast_size: int, forecast_size: int):\n",
    "        super().__init__()\n",
    "        frequency = np.append(np.zeros(1, dtype=float),\n",
    "                                        np.arange(harmonics, harmonics / 2 * forecast_size,\n",
    "                                                    dtype=float) / harmonics)[None, :]\n",
    "        backcast_grid = -2 * np.pi * (\n",
    "                np.arange(backcast_size, dtype=float)[:, None] / forecast_size) * frequency\n",
    "        forecast_grid = 2 * np.pi * (\n",
    "                np.arange(forecast_size, dtype=float)[:, None] / forecast_size) * frequency\n",
    "\n",
    "        backcast_cos_template = t.tensor(np.transpose(np.cos(backcast_grid)), dtype=t.float32)\n",
    "        backcast_sin_template = t.tensor(np.transpose(np.sin(backcast_grid)), dtype=t.float32)\n",
    "        backcast_template = t.cat([backcast_cos_template, backcast_sin_template], dim=0)\n",
    "\n",
    "        forecast_cos_template = t.tensor(np.transpose(np.cos(forecast_grid)), dtype=t.float32)\n",
    "        forecast_sin_template = t.tensor(np.transpose(np.sin(forecast_grid)), dtype=t.float32)\n",
    "        forecast_template = t.cat([forecast_cos_template, forecast_sin_template], dim=0)\n",
    "\n",
    "        self.backcast_basis = nn.Parameter(backcast_template, requires_grad=False)\n",
    "        self.forecast_basis = nn.Parameter(forecast_template, requires_grad=False)\n",
    "\n",
    "    def forward(self, theta: t.Tensor, insample_x_t: t.Tensor, outsample_x_t: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:\n",
    "        cut_point = self.forecast_basis.shape[0]\n",
    "        backcast = t.einsum('bp,pt->bt', theta[:, cut_point:], self.backcast_basis)\n",
    "        forecast = t.einsum('bp,pt->bt', theta[:, :cut_point], self.forecast_basis)\n",
    "        return backcast, forecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ExogenousBasisInterpretable(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, theta: t.Tensor, insample_x_t: t.Tensor, outsample_x_t: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:\n",
    "        backcast_basis = insample_x_t\n",
    "        forecast_basis = outsample_x_t\n",
    "\n",
    "        cut_point = forecast_basis.shape[1]\n",
    "        backcast = t.einsum('bp,bpt->bt', theta[:, cut_point:], backcast_basis)\n",
    "        forecast = t.einsum('bp,bpt->bt', theta[:, :cut_point], forecast_basis)\n",
    "        return backcast, forecast\n",
    "\n",
    "class ExogenousBasisWavenet(nn.Module):\n",
    "    def __init__(self, out_features, in_features, num_levels=4, kernel_size=3, dropout_prob=0):\n",
    "        super().__init__()\n",
    "        # Shape of (1, in_features, 1) to broadcast over b and t\n",
    "        self.weight = nn.Parameter(t.Tensor(1, in_features, 1), requires_grad=True)\n",
    "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(0.5))\n",
    "\n",
    "        padding = (kernel_size - 1) * (2**0)\n",
    "        input_layer = [nn.Conv1d(in_channels=in_features, out_channels=out_features,\n",
    "                                 kernel_size=kernel_size, padding=padding, dilation=2**0),\n",
    "                                 Chomp1d(padding),\n",
    "                                 nn.ReLU(),\n",
    "                                 nn.Dropout(dropout_prob)]\n",
    "        conv_layers = []\n",
    "        for i in range(1, num_levels):\n",
    "            dilation = 2**i\n",
    "            padding = (kernel_size - 1) * dilation\n",
    "            conv_layers.append(nn.Conv1d(in_channels=out_features, out_channels=out_features,\n",
    "                                         padding=padding, kernel_size=3, dilation=dilation))\n",
    "            conv_layers.append(Chomp1d(padding))\n",
    "            conv_layers.append(nn.ReLU())\n",
    "        conv_layers = input_layer + conv_layers\n",
    "\n",
    "        self.wavenet = nn.Sequential(*conv_layers)\n",
    "        \n",
    "    def transform(self, insample_x_t, outsample_x_t):\n",
    "        n_time_in = insample_x_t.shape[2]\n",
    "        \n",
    "        x_t = t.cat([insample_x_t, outsample_x_t], dim=2)\n",
    "        \n",
    "        x_t = x_t * self.weight # Element-wise multiplication, broadcasted on b and t. Weights used in L1 regularization\n",
    "        x_t = self.wavenet(x_t)[:]\n",
    "\n",
    "        backcast_basis = x_t[:,:, :n_time_in]\n",
    "        forecast_basis = x_t[:,:, n_time_in:]\n",
    "\n",
    "        return backcast_basis, forecast_basis\n",
    "\n",
    "    def forward(self, theta: t.Tensor, insample_x_t: t.Tensor, outsample_x_t: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:\n",
    "        backcast_basis, forecast_basis = self.transform(insample_x_t, outsample_x_t)\n",
    "\n",
    "        cut_point = forecast_basis.shape[1]\n",
    "        backcast = t.einsum('bp,bpt->bt', theta[:, cut_point:], backcast_basis)\n",
    "        forecast = t.einsum('bp,bpt->bt', theta[:, :cut_point], forecast_basis)\n",
    "        return backcast, forecast\n",
    "\n",
    "class ExogenousBasisTCN(nn.Module):\n",
    "    def __init__(self, out_features, in_features, num_levels = 4, kernel_size=2, dropout_prob=0):\n",
    "        super().__init__()\n",
    "        n_channels = num_levels * [out_features]\n",
    "        self.tcn = _TemporalConvNet(num_inputs=in_features, num_channels=n_channels, kernel_size=kernel_size, dropout=dropout_prob)\n",
    "        \n",
    "    def transform(self, insample_x_t, outsample_x_t):\n",
    "        n_time_in = insample_x_t.shape[2]\n",
    "        \n",
    "        x_t = t.cat([insample_x_t, outsample_x_t], dim=2)\n",
    "        \n",
    "        x_t = self.tcn(x_t)[:]\n",
    "        backcast_basis = x_t[:,:, :n_time_in]\n",
    "        forecast_basis = x_t[:,:, n_time_in:]\n",
    "\n",
    "        return backcast_basis, forecast_basis\n",
    "\n",
    "    def forward(self, theta: t.Tensor, insample_x_t: t.Tensor, outsample_x_t: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:\n",
    "        backcast_basis, forecast_basis = self.transform(insample_x_t, outsample_x_t)\n",
    "\n",
    "        cut_point = forecast_basis.shape[1]\n",
    "        backcast = t.einsum('bp,bpt->bt', theta[:, cut_point:], backcast_basis)\n",
    "        forecast = t.einsum('bp,bpt->bt', theta[:, :cut_point], forecast_basis)\n",
    "        return backcast, forecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def init_weights(module, initialization):\n",
    "    if type(module) == t.nn.Linear:\n",
    "        if initialization == 'orthogonal':\n",
    "            t.nn.init.orthogonal_(module.weight)\n",
    "        elif initialization == 'he_uniform':\n",
    "            t.nn.init.kaiming_uniform_(module.weight)\n",
    "        elif initialization == 'he_normal':\n",
    "            t.nn.init.kaiming_normal_(module.weight)\n",
    "        elif initialization == 'glorot_uniform':\n",
    "            t.nn.init.xavier_uniform_(module.weight)\n",
    "        elif initialization == 'glorot_normal':\n",
    "            t.nn.init.xavier_normal_(module.weight)\n",
    "        elif initialization == 'lecun_normal':\n",
    "            pass #t.nn.init.normal_(module.weight, 0.0, std=1/np.sqrt(module.weight.numel()))\n",
    "        else:\n",
    "            assert 1<0, f'Initialization {initialization} not found'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "ACTIVATIONS = ['ReLU',\n",
    "               'Softplus',\n",
    "               'Tanh',\n",
    "               'SELU',\n",
    "               'LeakyReLU',\n",
    "               'PReLU',\n",
    "               'Sigmoid']\n",
    "\n",
    "class _NBEATSBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    N-BEATS block which takes a basis function as an argument.\n",
    "    \"\"\"\n",
    "    def __init__(self, n_time_in: int, n_time_out: int, n_x: int,\n",
    "                 n_s: int, n_s_hidden: int, n_theta: int, n_theta_hidden: list,\n",
    "                 basis: nn.Module, \n",
    "                 n_layers: int,  batch_normalization: bool, dropout_prob: float, activation: str):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        if n_s == 0:\n",
    "            n_s_hidden = 0\n",
    "        n_theta_hidden = [n_time_in + (n_time_in+n_time_out)*n_x + n_s_hidden] + n_theta_hidden\n",
    "        \n",
    "        self.n_time_in = n_time_in\n",
    "        self.n_time_out = n_time_out\n",
    "        self.n_s = n_s\n",
    "        self.n_s_hidden = n_s_hidden\n",
    "        self.n_x = n_x\n",
    "        self.batch_normalization = batch_normalization\n",
    "        self.dropout_prob = dropout_prob\n",
    "        \n",
    "        assert activation in ACTIVATIONS, f'{activation} is not in {ACTIVATIONS}'\n",
    "        activ = getattr(nn, activation)()\n",
    "\n",
    "        hidden_layers = []\n",
    "        for i in range(n_layers):\n",
    "            hidden_layers.append(nn.Linear(in_features=n_theta_hidden[i], out_features=n_theta_hidden[i+1]))\n",
    "            hidden_layers.append(activ)\n",
    "\n",
    "            if self.batch_normalization:\n",
    "                hidden_layers.append(nn.BatchNorm1d(num_features=n_theta_hidden[i+1]))\n",
    "\n",
    "            if self.dropout_prob>0:\n",
    "                hidden_layers.append(nn.Dropout(p=self.dropout_prob))\n",
    "\n",
    "        output_layer = [nn.Linear(in_features=n_theta_hidden[-1], out_features=n_theta)]\n",
    "        layers = hidden_layers + output_layer\n",
    "\n",
    "        # n_s is computed with data, n_s_hidden is provided by user, if 0 no statics are used\n",
    "        if (self.n_s > 0) and (self.n_s_hidden > 0):\n",
    "            self.static_encoder = _StaticFeaturesEncoder(in_features=n_s, out_features=n_s_hidden)\n",
    "        self.layers = nn.Sequential(*layers)\n",
    "        self.basis = basis\n",
    "\n",
    "    def forward(self, insample_y: t.Tensor, insample_x_t: t.Tensor,\n",
    "                outsample_x_t: t.Tensor, x_s: t.Tensor) -> Tuple[t.Tensor, t.Tensor]:\n",
    "\n",
    "        batch_size = len(insample_y)\n",
    "        if self.n_x > 0:\n",
    "            insample_y = t.cat(( insample_y, insample_x_t.reshape(batch_size, -1) ), 1)\n",
    "            insample_y = t.cat(( insample_y, outsample_x_t.reshape(batch_size, -1) ), 1)\n",
    "        \n",
    "        # Static exogenous\n",
    "        if (self.n_s > 0) and (self.n_s_hidden > 0):\n",
    "            x_s = self.static_encoder(x_s)\n",
    "            insample_y = t.cat((insample_y, x_s), 1)\n",
    "\n",
    "        # Compute local projection weights and projection\n",
    "        theta = self.layers(insample_y)\n",
    "        backcast, forecast = self.basis(theta, insample_x_t, outsample_x_t)\n",
    "\n",
    "        return backcast, forecast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class _NBEATS(nn.Module):\n",
    "    \"\"\"\n",
    "    N-Beats Model.\n",
    "    \"\"\"\n",
    "    def __init__(self, \n",
    "                 n_time_in,\n",
    "                 n_time_out,\n",
    "                 n_s,\n",
    "                 n_x,\n",
    "                 n_s_hidden,\n",
    "                 n_x_hidden,\n",
    "                 n_polynomials,\n",
    "                 n_harmonics,\n",
    "                 stack_types: list,\n",
    "                 n_blocks: list,\n",
    "                 n_layers: list,\n",
    "                 n_theta_hidden: list,\n",
    "                 dropout_prob_theta,\n",
    "                 activation,\n",
    "                 initialization,\n",
    "                 batch_normalization,\n",
    "                 shared_weights):\n",
    "        super().__init__()\n",
    "\n",
    "        self.n_time_out = n_time_out\n",
    "\n",
    "        blocks = self.create_stack(stack_types=stack_types, \n",
    "                                   n_blocks=n_blocks,\n",
    "                                   n_time_in=n_time_in,\n",
    "                                   n_time_out=n_time_out,\n",
    "                                   n_x=n_x,\n",
    "                                   n_x_hidden=n_x_hidden,\n",
    "                                   n_s=n_s,\n",
    "                                   n_s_hidden=n_s_hidden,\n",
    "                                   n_layers=n_layers,\n",
    "                                   n_theta_hidden=n_theta_hidden,\n",
    "                                   batch_normalization=batch_normalization,\n",
    "                                   dropout_prob_theta=dropout_prob_theta,\n",
    "                                   activation=activation,\n",
    "                                   shared_weights=shared_weights,\n",
    "                                   n_polynomials=n_polynomials, \n",
    "                                   n_harmonics=n_harmonics,\n",
    "                                   initialization=initialization)\n",
    "        self.blocks = t.nn.ModuleList(blocks)\n",
    "\n",
    "    def create_stack(self, stack_types, n_blocks, \n",
    "                     n_time_in, n_time_out, \n",
    "                     n_x, n_x_hidden, n_s, n_s_hidden, \n",
    "                     n_layers, n_theta_hidden, batch_normalization, dropout_prob_theta, \n",
    "                     activation, shared_weights,\n",
    "                     n_polynomials, n_harmonics, initialization):                     \n",
    "\n",
    "        block_list = []\n",
    "        for i in range(len(stack_types)):\n",
    "            #print(f'| --  Stack {stack_types[i]} (#{i})')\n",
    "            for block_id in range(n_blocks[i]):\n",
    "                \n",
    "                # Batch norm only on first block\n",
    "                if (len(block_list)==0) and (batch_normalization):\n",
    "                    batch_normalization_block = True\n",
    "                else:\n",
    "                    batch_normalization_block = False\n",
    "\n",
    "                # Shared weights\n",
    "                if shared_weights and block_id>0:\n",
    "                    nbeats_block = block_list[-1]\n",
    "                else:\n",
    "                    if stack_types[i] == 'seasonality':\n",
    "                        n_theta = 4 * int(np.ceil(n_harmonics / 2 * n_time_out) - (n_harmonics - 1))\n",
    "                        basis = SeasonalityBasis(harmonics=n_harmonics,\n",
    "                                                 backcast_size=n_time_in,\n",
    "                                                 forecast_size=n_time_out)\n",
    "\n",
    "                    elif stack_types[i] == 'trend':\n",
    "                        n_theta = 2 * (n_polynomials + 1)\n",
    "                        basis = TrendBasis(degree_of_polynomial=n_polynomials,\n",
    "                                           backcast_size=n_time_in,\n",
    "                                           forecast_size=n_time_out)\n",
    "\n",
    "                    elif stack_types[i] == 'identity':\n",
    "                        n_theta = n_time_in + n_time_out\n",
    "                        basis = IdentityBasis(backcast_size=n_time_in,\n",
    "                                              forecast_size=n_time_out)                        \n",
    "\n",
    "                    elif stack_types[i] == 'exogenous':\n",
    "                        n_theta = 2 * n_x\n",
    "                        basis = ExogenousBasisInterpretable()\n",
    "\n",
    "                    elif stack_types[i] == 'exogenous_tcn':\n",
    "                        n_theta = 2 * n_x_hidden\n",
    "                        basis = ExogenousBasisTCN(n_x_hidden, n_x)\n",
    "\n",
    "                    elif stack_types[i] == 'exogenous_wavenet':\n",
    "                        n_theta = 2 * n_x_hidden\n",
    "                        basis = ExogenousBasisWavenet(n_x_hidden, n_x)\n",
    "\n",
    "                    else:\n",
    "                        assert 1<0, f'Block type not found!'\n",
    "\n",
    "                    nbeats_block = _NBEATSBlock(n_time_in=n_time_in,\n",
    "                                                n_time_out=n_time_out,\n",
    "                                                n_x=n_x,\n",
    "                                                n_s=n_s,\n",
    "                                                n_s_hidden=n_s_hidden,\n",
    "                                                n_theta=n_theta,\n",
    "                                                n_theta_hidden=n_theta_hidden[i],\n",
    "                                                basis=basis,\n",
    "                                                n_layers=n_layers[i],\n",
    "                                                batch_normalization=batch_normalization_block,\n",
    "                                                dropout_prob=dropout_prob_theta,\n",
    "                                                activation=activation)\n",
    "\n",
    "                # Select type of evaluation and apply it to all layers of block\n",
    "                init_function = partial(init_weights, initialization=initialization)                                             \n",
    "                nbeats_block.layers.apply(init_function)\n",
    "                #print(f'     | -- {nbeats_block}')\n",
    "                block_list.append(nbeats_block)\n",
    "        return block_list\n",
    "\n",
    "    def forward(self, S: t.Tensor, Y: t.Tensor, X: t.Tensor, \n",
    "                insample_mask: t.Tensor, outsample_mask: t.Tensor,\n",
    "                return_decomposition: bool=False):\n",
    "        \n",
    "        # insample\n",
    "        insample_y    = Y[:, :-self.n_time_out]\n",
    "        insample_x_t  = X[:, :, :-self.n_time_out]\n",
    "        insample_mask = insample_mask[:, :-self.n_time_out]\n",
    "        \n",
    "        # outsample\n",
    "        outsample_y   = Y[:, -self.n_time_out:]\n",
    "        outsample_x_t = X[:, :, -self.n_time_out:]\n",
    "        outsample_mask = outsample_mask[:, -self.n_time_out:]\n",
    "\n",
    "        if return_decomposition:\n",
    "            forecast, block_forecasts = self.forecast_decomposition(insample_y=insample_y, \n",
    "                                                                    insample_x_t=insample_x_t, \n",
    "                                                                    insample_mask=insample_mask,\n",
    "                                                                    outsample_x_t=outsample_x_t,\n",
    "                                                                    x_s=S)\n",
    "            return outsample_y, forecast, block_forecasts, outsample_mask\n",
    "        \n",
    "        else:\n",
    "            forecast = self.forecast(insample_y=insample_y,\n",
    "                                     insample_x_t=insample_x_t, \n",
    "                                     insample_mask=insample_mask,\n",
    "                                     outsample_x_t=outsample_x_t,\n",
    "                                     x_s=S)\n",
    "            return outsample_y, forecast, outsample_mask\n",
    "\n",
    "    def forecast(self, insample_y: t.Tensor, insample_x_t: t.Tensor, insample_mask: t.Tensor,\n",
    "                 outsample_x_t: t.Tensor, x_s: t.Tensor):\n",
    "\n",
    "        residuals = insample_y.flip(dims=(-1,))\n",
    "        insample_x_t = insample_x_t.flip(dims=(-1,))\n",
    "        insample_mask = insample_mask.flip(dims=(-1,))\n",
    "\n",
    "        forecast = insample_y[:, -1:] # Level with Naive1\n",
    "        for i, block in enumerate(self.blocks):\n",
    "            backcast, block_forecast = block(insample_y=residuals, insample_x_t=insample_x_t,\n",
    "                                             outsample_x_t=outsample_x_t, x_s=x_s)\n",
    "            residuals = (residuals - backcast) * insample_mask\n",
    "            forecast = forecast + block_forecast\n",
    "\n",
    "        return forecast\n",
    "\n",
    "    def forecast_decomposition(self, insample_y: t.Tensor, insample_x_t: t.Tensor, insample_mask: t.Tensor,\n",
    "                               outsample_x_t: t.Tensor, x_s: t.Tensor):\n",
    "\n",
    "        residuals = insample_y.flip(dims=(-1,))\n",
    "        insample_x_t = insample_x_t.flip(dims=(-1,))\n",
    "        insample_mask = insample_mask.flip(dims=(-1,))\n",
    "        \n",
    "        n_batch, n_channels, n_t = outsample_x_t.size(0), outsample_x_t.size(1), outsample_x_t.size(2)\n",
    "        \n",
    "        level = insample_y[:, -1:] # Level with Naive1\n",
    "        block_forecasts = [ level.repeat(1, n_t) ]\n",
    "                \n",
    "        forecast = level\n",
    "        for i, block in enumerate(self.blocks):\n",
    "            backcast, block_forecast = block(insample_y=residuals, insample_x_t=insample_x_t,\n",
    "                                             outsample_x_t=outsample_x_t, x_s=x_s)\n",
    "            residuals = (residuals - backcast) * insample_mask\n",
    "            forecast = forecast + block_forecast\n",
    "            block_forecasts.append(block_forecast)\n",
    "            \n",
    "        # (n_batch, n_blocks, n_t)\n",
    "        block_forecasts = t.stack(block_forecasts)\n",
    "        block_forecasts = block_forecasts.permute(1,0,2)\n",
    "\n",
    "        return forecast, block_forecasts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchinfo import summary\n",
    "\n",
    "model = _NBEATS(n_time_in=168,\n",
    "                n_time_out=24,\n",
    "                n_s=1,\n",
    "                n_x=10,\n",
    "                n_s_hidden=100,\n",
    "                n_x_hidden=30,\n",
    "                n_polynomials=2,\n",
    "                n_harmonics=4,\n",
    "                stack_types=['trend', 'seasonality', 'exogenous_wavenet'],\n",
    "                n_blocks=[1, 1, 1],\n",
    "                n_layers=[2, 2, 2],\n",
    "                n_theta_hidden=3 * [[128, 128]],\n",
    "                dropout_prob_theta=0,\n",
    "                activation='SELU',\n",
    "                initialization='lecun_normal',\n",
    "                batch_normalization=True,\n",
    "                shared_weights=True)\n",
    "\n",
    "# inputs: S, Y, X, insample_mask\n",
    "# S.shape (n_batch,n_s)\n",
    "# Y.shape (n_batch,n_time_in+n_time_out) \n",
    "# X.shape (n_batch,n_x,n_time_in+n_time_out)\n",
    "# insample_mask.shape (n_batch,n_time_in+n_time_out)\n",
    "# summary(model, input_size=[(256, 1), (256, 168+24), (256, 10, 168+24), (256, 168+24)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# N-BEATS model wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import time\n",
    "import random\n",
    "from copy import deepcopy\n",
    "from collections import defaultdict\n",
    "\n",
    "from torch import optim\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from nixtlats.losses.pytorch import (\n",
    "    MAPELoss, MASELoss, SMAPELoss, \n",
    "    MSELoss, MAELoss, PinballLoss\n",
    ")\n",
    "from nixtlats.losses.numpy import (\n",
    "    mae, mse, mape, \n",
    "    smape, rmse, pinball_loss\n",
    ")\n",
    "\n",
    "from nixtlats.data.tsdataset import WindowsDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class NBEATS(pl.LightningModule):\n",
    "    def __init__(self,\n",
    "                 n_time_in,\n",
    "                 n_time_out,\n",
    "                 n_x,\n",
    "                 n_x_hidden,\n",
    "                 n_s,\n",
    "                 n_s_hidden,\n",
    "                 shared_weights,\n",
    "                 activation,\n",
    "                 initialization,\n",
    "                 stack_types,\n",
    "                 n_blocks,\n",
    "                 n_layers,\n",
    "                 n_harmonics,\n",
    "                 n_polynomials,\n",
    "                 n_theta_hidden,\n",
    "                 batch_normalization,\n",
    "                 dropout_prob_theta,\n",
    "                 learning_rate,\n",
    "                 lr_decay,\n",
    "                 lr_decay_step_size,\n",
    "                 weight_decay,\n",
    "                 loss_train,\n",
    "                 loss_hypar,\n",
    "                 loss_valid,\n",
    "                 frequency,\n",
    "                 random_seed,\n",
    "                 seasonality):\n",
    "        super(NBEATS, self).__init__()\n",
    "        \"\"\"\n",
    "        N-BEATS model.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        # TODO: Fix parameters' documentation.\n",
    "        # TODO: Remove useless parameters (dropout_prob_exogenous).\n",
    "        n_time_in: int\n",
    "            Multiplier to get insample size.\n",
    "            Insample size = n_time_in * output_size\n",
    "        n_time_out: int\n",
    "            Forecast horizon.\n",
    "        shared_weights: bool\n",
    "            If True, repeats first block.\n",
    "        activation: str\n",
    "            Activation function.\n",
    "            An item from ['relu', 'softplus', 'tanh', 'selu', 'lrelu', 'prelu', 'sigmoid'].\n",
    "        initialization: str\n",
    "            Initialization function.\n",
    "            An item from ['orthogonal', 'he_uniform', 'glorot_uniform', 'glorot_normal', 'lecun_normal'].\n",
    "        stack_types: List[str]\n",
    "            List of stack types.\n",
    "            Subset from ['seasonality', 'trend', 'identity', 'exogenous', 'exogenous_tcn', 'exogenous_wavenet'].\n",
    "        n_blocks: List[int]\n",
    "            Number of blocks for each stack type.\n",
    "            Note that len(n_blocks) = len(stack_types).\n",
    "        n_layers: List[int]\n",
    "            Number of layers for each stack type.\n",
    "            Note that len(n_layers) = len(stack_types).\n",
    "        n_hidden: List[List[int]]\n",
    "            Structure of hidden layers for each stack type.\n",
    "            Each internal list should contain the number of units of each hidden layer.\n",
    "            Note that len(n_hidden) = len(stack_types).\n",
    "        n_harmonics: List[int]\n",
    "            Number of harmonic terms for each stack type.\n",
    "            Note that len(n_harmonics) = len(stack_types).\n",
    "        n_polynomials: List[int]\n",
    "            Number of polynomial terms for each stack type.\n",
    "            Note that len(n_polynomials) = len(stack_types).\n",
    "        exogenous_n_channels:\n",
    "            Exogenous channels for non-interpretable exogenous basis.\n",
    "        batch_normalization: bool\n",
    "            Whether perform batch normalization. \n",
    "        dropout_prob_theta: float\n",
    "            Float between (0, 1).\n",
    "            Dropout for Nbeats basis.\n",
    "        dropout_prob_exogenous: float\n",
    "            Float between (0, 1).\n",
    "            Dropout for exogenous basis.\n",
    "        x_s_n_hidden: int\n",
    "            Number of encoded static features to calculate.\n",
    "        learning_rate: float\n",
    "            Learning rate between (0, 1).\n",
    "        lr_decay: float\n",
    "            Decreasing multiplier for the learning rate.\n",
    "        lr_decay_step_size: int\n",
    "            Steps between each lerning rate decay.\n",
    "        weight_decay: float\n",
    "            L2 penalty for optimizer.\n",
    "        loss_train: str\n",
    "            Loss to optimize.\n",
    "            An item from ['MAPE', 'MASE', 'SMAPE', 'MSE', 'MAE', 'PINBALL', 'PINBALL2'].\n",
    "        loss_hypar:\n",
    "            Hyperparameter for chosen loss.\n",
    "        loss_valid:\n",
    "            Validation loss.\n",
    "            An item from ['MAPE', 'MASE', 'SMAPE', 'RMSE', 'MAE', 'PINBALL'].\n",
    "        frequency: str\n",
    "            Time series frequency.\n",
    "        random_seed: int\n",
    "            random_seed for pseudo random pytorch initializer and\n",
    "            numpy random generator.\n",
    "        seasonality: int\n",
    "            Time series seasonality.\n",
    "            Usually 7 for daily data, 12 for monthly data and 4 for weekly data.\n",
    "        \"\"\"\n",
    "\n",
    "        if activation == 'SELU': initialization = 'lecun_normal'\n",
    "\n",
    "        #------------------------ Model Attributes ------------------------#\n",
    "        # Architecture parameters\n",
    "        self.n_time_in = n_time_in\n",
    "        self.n_time_out = n_time_out\n",
    "        self.n_x = n_x\n",
    "        self.n_x_hidden = n_x_hidden\n",
    "        self.n_s = n_s\n",
    "        self.n_s_hidden = n_s_hidden\n",
    "        self.shared_weights = shared_weights\n",
    "        self.activation = activation\n",
    "        self.initialization = initialization\n",
    "        self.stack_types = stack_types\n",
    "        self.n_blocks = n_blocks\n",
    "        self.n_layers = n_layers\n",
    "        self.n_harmonics = n_harmonics\n",
    "        self.n_polynomials = n_polynomials\n",
    "        self.n_theta_hidden = n_theta_hidden\n",
    "\n",
    "        # Loss functions\n",
    "        self.loss_train = loss_train\n",
    "        self.loss_hypar = loss_hypar\n",
    "        self.loss_valid = loss_valid\n",
    "        self.loss_fn_train = LossFunction(loss_train, \n",
    "                                          seasonality=self.loss_hypar)\n",
    "        self.loss_fn_valid = LossFunction(loss_valid,\n",
    "                                          seasonality=self.loss_hypar)\n",
    "        \n",
    "        # Regularization and optimization parameters\n",
    "        self.batch_normalization = batch_normalization\n",
    "        self.dropout_prob_theta = dropout_prob_theta        \n",
    "        self.learning_rate = learning_rate\n",
    "        self.lr_decay = lr_decay\n",
    "        self.weight_decay = weight_decay\n",
    "        self.lr_decay_step_size = lr_decay_step_size\n",
    "        self.random_seed = random_seed\n",
    "\n",
    "        # Data parameters\n",
    "        self.frequency = frequency\n",
    "        self.seasonality = seasonality\n",
    "        self.return_decomposition = False\n",
    "\n",
    "        self.model = _NBEATS(n_time_in=self.n_time_in,\n",
    "                             n_time_out=self.n_time_out,\n",
    "                             n_s=self.n_s,\n",
    "                             n_x=self.n_x,\n",
    "                             n_s_hidden=self.n_s_hidden,\n",
    "                             n_x_hidden=self.n_x_hidden,\n",
    "                             n_polynomials=self.n_polynomials,\n",
    "                             n_harmonics=self.n_harmonics,\n",
    "                             stack_types=self.stack_types,\n",
    "                             n_blocks=self.n_blocks,\n",
    "                             n_layers=self.n_layers,\n",
    "                             n_theta_hidden=self.n_theta_hidden,\n",
    "                             dropout_prob_theta=self.dropout_prob_theta,\n",
    "                             activation=self.activation,\n",
    "                             initialization=self.initialization,\n",
    "                             batch_normalization=self.batch_normalization,\n",
    "                             shared_weights=self.shared_weights)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        S = batch['S']\n",
    "        Y = batch['Y']\n",
    "        X = batch['X']\n",
    "        sample_mask = batch['sample_mask']\n",
    "        available_mask = batch['available_mask']\n",
    "\n",
    "        outsample_y, forecast, outsample_mask = self.model(S=S, Y=Y, X=X,\n",
    "                                                           insample_mask=available_mask,\n",
    "                                                           outsample_mask=sample_mask,\n",
    "                                                           return_decomposition=False)\n",
    "\n",
    "        loss = self.loss_fn_train(y=outsample_y,\n",
    "                                  y_hat=forecast,\n",
    "                                  mask=outsample_mask,\n",
    "                                  y_insample=Y)\n",
    "\n",
    "        self.log('train_loss', loss, prog_bar=True, on_epoch=True)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, idx):\n",
    "        S = batch['S']\n",
    "        Y = batch['Y']\n",
    "        X = batch['X']\n",
    "        sample_mask = batch['sample_mask']\n",
    "        available_mask = batch['available_mask']\n",
    "\n",
    "        outsample_y, forecast, outsample_mask = self.model(S=S, Y=Y, X=X,\n",
    "                                                           insample_mask=available_mask,\n",
    "                                                           outsample_mask=sample_mask,\n",
    "                                                           return_decomposition=False)\n",
    "\n",
    "        loss = self.loss_fn_valid(y=outsample_y,\n",
    "                                  y_hat=forecast,\n",
    "                                  mask=outsample_mask,\n",
    "                                  y_insample=Y)\n",
    "\n",
    "        self.log('val_loss', loss, prog_bar=True)\n",
    "        \n",
    "        return loss\n",
    "\n",
    "    def on_fit_start(self):\n",
    "        t.manual_seed(self.random_seed)\n",
    "        np.random.seed(self.random_seed)\n",
    "        random.seed(self.random_seed) #TODO: interaccion rara con window_sampling de validacion\n",
    "\n",
    "    def forward(self, batch):\n",
    "        S = batch['S']\n",
    "        Y = batch['Y']\n",
    "        X = batch['X']\n",
    "        sample_mask = batch['sample_mask']\n",
    "        available_mask = batch['available_mask']\n",
    "\n",
    "        if self.return_decomposition:\n",
    "            outsample_y, forecast, block_forecast, outsample_mask = self.model(S=S, Y=Y, X=X,\n",
    "                                                                     insample_mask=available_mask,\n",
    "                                                                     outsample_mask=sample_mask,\n",
    "                                                                     return_decomposition=True)\n",
    "            return outsample_y, forecast, block_forecast, outsample_mask\n",
    "\n",
    "        outsample_y, forecast, outsample_mask = self.model(S=S, Y=Y, X=X,\n",
    "                                                           insample_mask=available_mask,\n",
    "                                                           outsample_mask=sample_mask,\n",
    "                                                           return_decomposition=False)\n",
    "        return outsample_y, forecast, outsample_mask\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = optim.Adam(self.model.parameters(),\n",
    "                               lr=self.learning_rate, \n",
    "                               weight_decay=self.weight_decay)\n",
    "        \n",
    "        lr_scheduler = optim.lr_scheduler.StepLR(optimizer, \n",
    "                                                 step_size=self.lr_decay_step_size, \n",
    "                                                 gamma=self.lr_decay)\n",
    "\n",
    "        return {'optimizer': optimizer, 'lr_scheduler': lr_scheduler}\n",
    "     "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# N-BEATS Usage Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 56356 characters]\n",
      "text/plain": [
       "<Figure size 1080x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import multiprocessing\n",
    "\n",
    "import pandas as pd\n",
    "from nixtlats.data.datasets.epf import EPF#, EPFInfo\n",
    "from nixtlats.data.tsloader import TimeSeriesLoader\n",
    "\n",
    "import pylab as plt\n",
    "from pylab import rcParams\n",
    "plt.style.use('seaborn-whitegrid')\n",
    "plt.rcParams['font.family'] = 'serif'\n",
    "\n",
    "FONTSIZE = 19\n",
    "\n",
    "# Load and plot data\n",
    "Y_df, X_df, S_df = EPF.load_groups(directory='./data', groups=['NP'])\n",
    "\n",
    "fig = plt.figure(figsize=(15, 6))\n",
    "plt.plot(Y_df.ds, Y_df.y.values, color='#628793', linewidth=0.4)\n",
    "plt.ylabel('Price [EUR/MWh]', fontsize=19)\n",
    "plt.xlabel('Date', fontsize=15)\n",
    "# plt.savefig('./results/NP.png', bbox_inches = 'tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Declare Model and Data Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=================================================================\n",
      "model                                                      nbeats\n",
      "mode                                                       simple\n",
      "activation                                                   SELU\n",
      "n_time_in                                                     168\n",
      "n_time_out                                                    168\n",
      "n_x_hidden                                                      8\n",
      "n_s_hidden                                                      0\n",
      "stack_types               [trend, seasonality, exogenous_wavenet]\n",
      "n_blocks                                                [1, 1, 1]\n",
      "n_layers                                                [2, 2, 2]\n",
      "n_hidden                                                      128\n",
      "shared_weights                                              False\n",
      "n_harmonics                                                     2\n",
      "n_polynomials                                                   4\n",
      "initialization                                       lecun_normal\n",
      "learning_rate                                              0.0005\n",
      "batch_size                                                    256\n",
      "lr_decay                                                      0.5\n",
      "lr_decay_step_size                                              2\n",
      "max_epochs                                                      1\n",
      "max_steps                                                      20\n",
      "early_stop_patience                                            20\n",
      "eval_freq                                                     500\n",
      "batch_normalization                                         False\n",
      "dropout_prob_theta                                              0\n",
      "dropout_prob_exogenous                                          0\n",
      "l1_theta                                                        0\n",
      "weight_decay                                                6e-05\n",
      "loss_train                                                    MAE\n",
      "loss_hypar                                                    0.5\n",
      "loss_valid                                                    MAE\n",
      "random_seed                                                     1\n",
      "len_sample_chunks                                            None\n",
      "idx_to_sample_freq                                              1\n",
      "val_idx_to_sample_freq                                        168\n",
      "n_val_weeks                                                    52\n",
      "window_sampling_limit                                      500000\n",
      "normalizer_y                                                 None\n",
      "normalizer_x                                               median\n",
      "complete_inputs                                             False\n",
      "frequency                                                       H\n",
      "seasonality                                                    24\n",
      "dtype: object\n",
      "=================================================================\n"
     ]
    }
   ],
   "source": [
    "# Architecture parameters\n",
    "mc = {}\n",
    "mc['model'] = 'nbeats'\n",
    "mc['mode'] = 'simple'\n",
    "mc['activation'] = 'SELU'\n",
    "\n",
    "mc['n_time_in'] = 24*7\n",
    "mc['n_time_out'] = 24*7\n",
    "mc['n_x_hidden'] = 8\n",
    "mc['n_s_hidden'] = 0\n",
    "\n",
    "# mc['input_size_multiplier'] = 7\n",
    "# mc['output_size'] = 24\n",
    "\n",
    "mc['stack_types'] = ['trend', 'seasonality', 'exogenous_wavenet']\n",
    "mc['n_blocks'] = [1, 1, 1]\n",
    "mc['n_layers'] = [2, 2, 2]\n",
    "# mc['stack_types'] = ['trend', 'seasonality']\n",
    "# mc['n_blocks'] = [1, 1]\n",
    "# mc['n_layers'] = [2, 2]\n",
    "\n",
    "mc['n_hidden'] = 128\n",
    "mc['shared_weights'] = False\n",
    "mc['n_harmonics'] = 4\n",
    "mc['n_polynomials'] = 2\n",
    "\n",
    "# Optimization and regularization parameters\n",
    "mc['initialization'] = 'lecun_normal'\n",
    "mc['learning_rate'] = 0.0007\n",
    "mc['batch_size'] = 128\n",
    "mc['lr_decay'] = 0.5\n",
    "mc['lr_decay_step_size'] = 2\n",
    "mc['max_epochs'] = 1#_000\n",
    "mc['max_steps'] = 20#_000\n",
    "mc['early_stop_patience'] = 20\n",
    "mc['eval_freq'] = 500\n",
    "mc['batch_normalization'] = False\n",
    "mc['dropout_prob_theta'] = 0.51\n",
    "mc['dropout_prob_exogenous'] = 0.44\n",
    "mc['l1_theta'] = 0\n",
    "mc['weight_decay'] = 0\n",
    "mc['loss_train'] = 'MAE'\n",
    "mc['loss_hypar'] = 0.5\n",
    "mc['loss_valid'] = mc['loss_train']\n",
    "mc['random_seed'] = 1\n",
    "\n",
    "# Data Parameters\n",
    "mc['len_sample_chunks'] = None\n",
    "mc['idx_to_sample_freq'] = 1\n",
    "mc['val_idx_to_sample_freq'] = 24 * 7\n",
    "mc['n_val_weeks'] = 52\n",
    "mc['window_sampling_limit'] = 500_000\n",
    "mc['normalizer_y'] = None\n",
    "mc['normalizer_x'] = 'median'\n",
    "mc['complete_inputs'] = False\n",
    "mc['frequency'] = 'H'\n",
    "mc['seasonality'] = 24\n",
    "\n",
    "# # Within decomposition\n",
    "mc['learning_rate'] = 0.0005\n",
    "mc['batch_size'] = 256\n",
    "mc['weight_decay'] = 0.00006\n",
    "mc['n_harmonics'] = 2\n",
    "mc['n_polynomials'] = 4\n",
    "mc['dropout_prob_theta'] = 0\n",
    "mc['dropout_prob_exogenous'] = 0\n",
    "\n",
    "print(65*'=')\n",
    "print(pd.Series(mc))\n",
    "print(65*'=')\n",
    "\n",
    "mc['n_theta_hidden'] = len(mc['stack_types']) * [ [int(mc['n_hidden']), int(mc['n_hidden'])] ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate Loaders and Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/statsmodels/tools/_testing.py:19: FutureWarning: pandas.util.testing is deprecated. Use the functions in the public API at pandas.testing instead.\n",
      "  import pandas.util.testing as tm\n",
      "INFO:root:Train Validation splits\n",
      "\n",
      "INFO:root:                              ds                    \n",
      "                             min                 max\n",
      "unique_id sample_mask                               \n",
      "NP        0           2016-03-08 2018-12-24 23:00:00\n",
      "          1           2013-01-01 2016-03-07 23:00:00\n",
      "INFO:root:\n",
      "Total data \t\t\t52416 time stamps \n",
      "Available percentage=100.0, \t52416 time stamps \n",
      "Insample  percentage=53.21, \t27888 time stamps \n",
      "Outsample percentage=46.79, \t24528 time stamps \n",
      "\n",
      "INFO:root:Train Validation splits\n",
      "\n",
      "INFO:root:                              ds                    \n",
      "                             min                 max\n",
      "unique_id sample_mask                               \n",
      "NP        0           2013-01-01 2018-12-24 23:00:00\n",
      "          1           2016-03-08 2016-12-26 23:00:00\n",
      "INFO:root:\n",
      "Total data \t\t\t52416 time stamps \n",
      "Available percentage=100.0, \t52416 time stamps \n",
      "Insample  percentage=13.46, \t7056 time stamps \n",
      "Outsample percentage=86.54, \t45360 time stamps \n",
      "\n",
      "INFO:root:Train Validation splits\n",
      "\n",
      "INFO:root:                              ds                    \n",
      "                             min                 max\n",
      "unique_id sample_mask                               \n",
      "NP        0           2013-01-01 2016-12-26 23:00:00\n",
      "          1           2016-12-27 2018-12-24 23:00:00\n",
      "INFO:root:\n",
      "Total data \t\t\t52416 time stamps \n",
      "Available percentage=100.0, \t52416 time stamps \n",
      "Insample  percentage=33.33, \t17472 time stamps \n",
      "Outsample percentage=66.67, \t34944 time stamps \n",
      "\n"
     ]
    }
   ],
   "source": [
    "from nixtlats.experiments.utils import create_datasets\n",
    "\n",
    "train_dataset, val_dataset, test_dataset, scaler_y = create_datasets(mc=mc,\n",
    "                                                                     S_df=S_df, Y_df=Y_df, X_df=X_df,\n",
    "                                                                     f_cols=['Exogenous1', 'Exogenous2'],\n",
    "                                                                     ds_in_val=294*24,\n",
    "                                                                     ds_in_test=728*24,\n",
    "                                                                     n_uids=None, n_val_windows=None,freq=None, \n",
    "                                                                     is_val_random=False)\n",
    "\n",
    "train_loader = TimeSeriesLoader(dataset=train_dataset,\n",
    "                                batch_size=int(mc['batch_size']),\n",
    "                                #num_workers=int(min(multiprocessing.cpu_count(), 3)),\n",
    "                                shuffle=True)\n",
    "\n",
    "val_loader = TimeSeriesLoader(dataset=val_dataset,\n",
    "                              batch_size=int(mc['batch_size']),\n",
    "                              #num_workers=int(min(multiprocessing.cpu_count(), 3)),\n",
    "                              shuffle=False)\n",
    "\n",
    "test_loader = TimeSeriesLoader(dataset=test_dataset,\n",
    "                               batch_size=int(mc['batch_size']),\n",
    "                               #num_workers=int(min(multiprocessing.cpu_count(), 3)),\n",
    "                               shuffle=False)\n",
    "\n",
    "mc['n_x'], mc['n_s'] = train_dataset.get_n_variables()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NBEATS(n_time_in=int(mc['n_time_in']),\n",
    "               n_time_out=int(mc['n_time_out']),\n",
    "               n_x=mc['n_x'],\n",
    "               n_s=mc['n_s'],\n",
    "               n_s_hidden=int(mc['n_s_hidden']),\n",
    "               n_x_hidden=int(mc['n_x_hidden']),\n",
    "               shared_weights=mc['shared_weights'],\n",
    "               initialization=mc['initialization'],\n",
    "               activation=mc['activation'],\n",
    "               stack_types=mc['stack_types'],\n",
    "               n_blocks=mc['n_blocks'],\n",
    "               n_layers=mc['n_layers'],\n",
    "               n_theta_hidden=mc['n_theta_hidden'],\n",
    "               n_harmonics=int(mc['n_harmonics']),\n",
    "               n_polynomials=int(mc['n_polynomials']),\n",
    "               batch_normalization = mc['batch_normalization'],\n",
    "               dropout_prob_theta=mc['dropout_prob_theta'],\n",
    "               learning_rate=float(mc['learning_rate']),\n",
    "               lr_decay=float(mc['lr_decay']),\n",
    "               lr_decay_step_size=float(mc['lr_decay_step_size']),\n",
    "               weight_decay=mc['weight_decay'],\n",
    "               loss_train=mc['loss_train'],\n",
    "               loss_hypar=float(mc['loss_hypar']),\n",
    "               loss_valid=mc['loss_valid'],\n",
    "               frequency=mc['frequency'],\n",
    "               seasonality=int(mc['seasonality']),\n",
    "               random_seed=int(mc['random_seed']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/accelerator_connector.py:597: UserWarning: GPU available but not used. Set the gpus flag in your trainer `Trainer(gpus=1)` or script `--gpus=1`.\n",
      "  \"GPU available but not used. Set the gpus flag in your trainer\"\n",
      "\n",
      "  | Name  | Type    | Params\n",
      "----------------------------------\n",
      "0 | model | _NBEATS | 1.6 M \n",
      "----------------------------------\n",
      "1.5 M     Trainable params\n",
      "113 K     Non-trainable params\n",
      "1.6 M     Total params\n",
      "6.436     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9118f3b66cb84307b1e79bc1745d6eee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/data_loading.py:103: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 4 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  f'The dataloader, {name}, does not have many workers which may be a bottleneck.'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/data_loading.py:103: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 4 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  f'The dataloader, {name}, does not have many workers which may be a bottleneck.'\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0777009ed9441b2bb937a53d55e80b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54b1686a18744855b855986d1df1208f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py:611: LightningDeprecationWarning: Relying on `self.log('val_loss', ...)` to set the ModelCheckpoint monitor is deprecated in v1.2 and will be removed in v1.4. Please, create your own `mc = ModelCheckpoint(monitor='your_monitor')` and use it as `Trainer(callbacks=[mc])`.\n",
      "  \"Relying on `self.log('val_loss', ...)` to set the ModelCheckpoint monitor is deprecated in v1.2\"\n"
     ]
    }
   ],
   "source": [
    "from pytorch_lightning.callbacks import EarlyStopping\n",
    "\n",
    "early_stopping = EarlyStopping(monitor=\"val_loss\", \n",
    "                               min_delta=1e-4, \n",
    "                               patience=mc['early_stop_patience'],\n",
    "                               verbose=False,\n",
    "                               mode=\"min\")\n",
    "\n",
    "trainer = pl.Trainer(max_epochs=mc['max_epochs'], \n",
    "                     max_steps=mc['max_steps'],\n",
    "                     gradient_clip_val=1.0,\n",
    "                     progress_bar_refresh_rate=10, \n",
    "                     log_every_n_steps=500, \n",
    "                     check_val_every_n_epoch=1,\n",
    "                     callbacks=[early_stopping])\n",
    "\n",
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/data_loading.py:103: UserWarning: The dataloader, predict dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 4 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  f'The dataloader, {name}, does not have many workers which may be a bottleneck.'\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0493c9bfed3446f58c26c89f8c308626",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Predicting', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "outputs[0][0].shape torch.Size([42, 168])\n",
      "outputs[0][1].shape torch.Size([42, 168])\n",
      "outputs[0][2].shape torch.Size([42, 4, 168])\n"
     ]
    }
   ],
   "source": [
    "model.return_decomposition = True\n",
    "outputs = trainer.predict(model, val_loader)\n",
    "\n",
    "print(\"outputs[0][0].shape\", outputs[0][0].shape)\n",
    "print(\"outputs[0][1].shape\", outputs[0][1].shape)\n",
    "print(\"outputs[0][2].shape\", outputs[0][2].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
