# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Temporal Fusion Transformer Model.

Contains the full TFT architecture and associated components. Defines functions
for training, evaluation and prediction using simple Pandas Dataframe inputs.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import gc
import json
import os
import shutil

import data_formatters.base
import libs.utils as utils
import numpy as np
import pandas as pd
import tensorflow as tf

# Layer definitions.
concat = tf.keras.backend.concatenate
stack = tf.keras.backend.stack
K = tf.keras.backend
Add = tf.keras.layers.Add
LayerNorm = tf.keras.layers.LayerNormalization
Dense = tf.keras.layers.Dense
Multiply = tf.keras.layers.Multiply
Dropout = tf.keras.layers.Dropout
Activation = tf.keras.layers.Activation
Lambda = tf.keras.layers.Lambda

# Default input types.
InputTypes = data_formatters.base.InputTypes


# Layer utility functions.
def linear_layer(size, activation=None, use_time_distributed=False, use_bias=True):
    """Returns simple Keras linear layer.

    Args:
      size: Output size
      activation: Activation function to apply if required
      use_time_distributed: Whether to apply layer across time
      use_bias: Whether bias should be included in layer
    """
    linear = tf.keras.layers.Dense(size, activation=activation, use_bias=use_bias)
    if use_time_distributed:
        linear = tf.keras.layers.TimeDistributed(linear)
    return linear


def apply_mlp(
    inputs, hidden_size, output_size, output_activation=None, hidden_activation="tanh", use_time_distributed=False
):
    """Applies simple feed-forward network to an input.

    Args:
      inputs: MLP inputs
      hidden_size: Hidden state size
      output_size: Output size of MLP
      output_activation: Activation function to apply on output
      hidden_activation: Activation function to apply on input
      use_time_distributed: Whether to apply across time

    Returns:
      Tensor for MLP outputs.
    """
    if use_time_distributed:
        hidden = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(hidden_size, activation=hidden_activation))(
            inputs
        )
        return tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(output_size, activation=output_activation))(hidden)
    else:
        hidden = tf.keras.layers.Dense(hidden_size, activation=hidden_activation)(inputs)
        return tf.keras.layers.Dense(output_size, activation=output_activation)(hidden)


def apply_gating_layer(x, hidden_layer_size, dropout_rate=None, use_time_distributed=True, activation=None):
    """Applies a Gated Linear Unit (GLU) to an input.

    Args:
      x: Input to gating layer
      hidden_layer_size: Dimension of GLU
      dropout_rate: Dropout rate to apply if any
      use_time_distributed: Whether to apply across time
      activation: Activation function to apply to the linear feature transform if
        necessary

    Returns:
      Tuple of tensors for: (GLU output, gate)
    """

    if dropout_rate is not None:
        x = tf.keras.layers.Dropout(dropout_rate)(x)

    if use_time_distributed:
        activation_layer = tf.keras.layers.TimeDistributed(
            tf.keras.layers.Dense(hidden_layer_size, activation=activation)
        )(x)
        gated_layer = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(hidden_layer_size, activation="sigmoid"))(x)
    else:
        activation_layer = tf.keras.layers.Dense(hidden_layer_size, activation=activation)(x)
        gated_layer = tf.keras.layers.Dense(hidden_layer_size, activation="sigmoid")(x)

    return tf.keras.layers.Multiply()([activation_layer, gated_layer]), gated_layer


def add_and_norm(x_list):
    """Applies skip connection followed by layer normalisation.

    Args:
      x_list: List of inputs to sum for skip connection

    Returns:
      Tensor output from layer.
    """
    tmp = Add()(x_list)
    tmp = LayerNorm()(tmp)
    return tmp


def gated_residual_network(
    x,
    hidden_layer_size,
    output_size=None,
    dropout_rate=None,
    use_time_distributed=True,
    additional_context=None,
    return_gate=False,
):
    """Applies the gated residual network (GRN) as defined in paper.

    Args:
      x: Network inputs
      hidden_layer_size: Internal state size
      output_size: Size of output layer
      dropout_rate: Dropout rate if dropout is applied
      use_time_distributed: Whether to apply network across time dimension
      additional_context: Additional context vector to use if relevant
      return_gate: Whether to return GLU gate for diagnostic purposes

    Returns:
      Tuple of tensors for: (GRN output, GLU gate)
    """

    # Setup skip connection
    if output_size is None:
        output_size = hidden_layer_size
        skip = x
    else:
        linear = Dense(output_size)
        if use_time_distributed:
            linear = tf.keras.layers.TimeDistributed(linear)
        skip = linear(x)

    # Apply feedforward network
    hidden = linear_layer(hidden_layer_size, activation=None, use_time_distributed=use_time_distributed)(x)
    if additional_context is not None:
        hidden = hidden + linear_layer(
            hidden_layer_size, activation=None, use_time_distributed=use_time_distributed, use_bias=False
        )(additional_context)
    hidden = tf.keras.layers.Activation("elu")(hidden)
    hidden = linear_layer(hidden_layer_size, activation=None, use_time_distributed=use_time_distributed)(hidden)

    gating_layer, gate = apply_gating_layer(
        hidden, output_size, dropout_rate=dropout_rate, use_time_distributed=use_time_distributed, activation=None
    )

    if return_gate:
        return add_and_norm([skip, gating_layer]), gate
    else:
        return add_and_norm([skip, gating_layer])


# Attention Components.
def get_decoder_mask(self_attn_inputs):
    """Returns causal mask to apply for self-attention layer.

    Args:
      self_attn_inputs: Inputs to self attention layer to determine mask shape
    """
    len_s = tf.shape(self_attn_inputs)[1]
    bs = tf.shape(self_attn_inputs)[:1]
    mask = K.cumsum(tf.eye(len_s, batch_shape=bs), 1)
    return mask


class ScaledDotProductAttention:
    """Defines scaled dot product attention layer.

    Attributes:
      dropout: Dropout rate to use
      activation: Normalisation function for scaled dot product attention (e.g.
        softmax by default)
    """

    def __init__(self, attn_dropout=0.0):
        self.dropout = Dropout(attn_dropout)
        self.activation = Activation("softmax")

    def __call__(self, q, k, v, mask):
        """Applies scaled dot product attention.

        Args:
          q: Queries
          k: Keys
          v: Values
          mask: Masking if required -- sets softmax to very large value

        Returns:
          Tuple of (layer outputs, attention weights)
        """
        temper = tf.sqrt(tf.cast(tf.shape(k)[-1], dtype="float32"))
        attn = Lambda(lambda x: K.batch_dot(x[0], x[1], axes=[2, 2]) / temper)([q, k])  # shape=(batch, q, k)
        if mask is not None:
            mmask = Lambda(lambda x: (-1e9) * (1.0 - K.cast(x, "float32")))(mask)  # setting to infinity
            attn = Add()([attn, mmask])
        attn = self.activation(attn)
        attn = self.dropout(attn)
        output = Lambda(lambda x: K.batch_dot(x[0], x[1]))([attn, v])
        return output, attn


class InterpretableMultiHeadAttention:
    """Defines interpretable multi-head attention layer.

    Attributes:
      n_head: Number of heads
      d_k: Key/query dimensionality per head
      d_v: Value dimensionality
      dropout: Dropout rate to apply
      qs_layers: List of queries across heads
      ks_layers: List of keys across heads
      vs_layers: List of values across heads
      attention: Scaled dot product attention layer
      w_o: Output weight matrix to project internal state to the original TFT
        state size
    """

    def __init__(self, n_head, d_model, dropout):
        """Initialises layer.

        Args:
          n_head: Number of heads
          d_model: TFT state dimensionality
          dropout: Dropout discard rate
        """

        self.n_head = n_head
        self.d_k = self.d_v = d_k = d_v = d_model // n_head
        self.dropout = dropout

        self.qs_layers = []
        self.ks_layers = []
        self.vs_layers = []

        # Use same value layer to facilitate interp
        vs_layer = Dense(d_v, use_bias=False)

        for _ in range(n_head):
            self.qs_layers.append(Dense(d_k, use_bias=False))
            self.ks_layers.append(Dense(d_k, use_bias=False))
            self.vs_layers.append(vs_layer)  # use same vs_layer

        self.attention = ScaledDotProductAttention()
        self.w_o = Dense(d_model, use_bias=False)

    def __call__(self, q, k, v, mask=None):
        """Applies interpretable multihead attention.

        Using T to denote the number of time steps fed into the transformer.

        Args:
          q: Query tensor of shape=(?, T, d_model)
          k: Key of shape=(?, T, d_model)
          v: Values of shape=(?, T, d_model)
          mask: Masking if required with shape=(?, T, T)

        Returns:
          Tuple of (layer outputs, attention weights)
        """
        n_head = self.n_head

        heads = []
        attns = []
        for i in range(n_head):
            qs = self.qs_layers[i](q)
            ks = self.ks_layers[i](k)
            vs = self.vs_layers[i](v)
            head, attn = self.attention(qs, ks, vs, mask)

            head_dropout = Dropout(self.dropout)(head)
            heads.append(head_dropout)
            attns.append(attn)
        head = K.stack(heads) if n_head > 1 else heads[0]
        attn = K.stack(attns)

        outputs = K.mean(head, axis=0) if n_head > 1 else head
        outputs = self.w_o(outputs)
        outputs = Dropout(self.dropout)(outputs)  # output dropout

        return outputs, attn


class TFTDataCache:
    """Caches data for the TFT."""

    _data_cache = {}

    @classmethod
    def update(cls, data, key):
        """Updates cached data.

        Args:
          data: Source to update
          key: Key to dictionary location
        """
        cls._data_cache[key] = data

    @classmethod
    def get(cls, key):
        """Returns data stored at key location."""
        return cls._data_cache[key].copy()

    @classmethod
    def contains(cls, key):
        """Retuns boolean indicating whether key is present in cache."""

        return key in cls._data_cache


# TFT model definitions.
class TemporalFusionTransformer:
    """Defines Temporal Fusion Transformer.

    Attributes:
      name: Name of model
      time_steps: Total number of input time steps per forecast date (i.e. Width
        of Temporal fusion decoder N)
      input_size: Total number of inputs
      output_size: Total number of outputs
      category_counts: Number of categories per categorical variable
      n_multiprocessing_workers: Number of workers to use for parallel
        computations
      column_definition: List of tuples of (string, DataType, InputType) that
        define each column
      quantiles: Quantiles to forecast for TFT
      use_cudnn: Whether to use Keras CuDNNLSTM or standard LSTM layers
      hidden_layer_size: Internal state size of TFT
      dropout_rate: Dropout discard rate
      max_gradient_norm: Maximum norm for gradient clipping
      learning_rate: Initial learning rate of ADAM optimizer
      minibatch_size: Size of minibatches for training
      num_epochs: Maximum number of epochs for training
      early_stopping_patience: Maximum number of iterations of non-improvement
        before early stopping kicks in
      num_encoder_steps: Size of LSTM encoder -- i.e. number of past time steps
        before forecast date to use
      num_stacks: Number of self-attention layers to apply (default is 1 for basic
        TFT)
      num_heads: Number of heads for interpretable mulit-head attention
      model: Keras model for TFT
    """

    def __init__(self, raw_params, use_cudnn=False):
        """Builds TFT from parameters.

        Args:
          raw_params: Parameters to define TFT
          use_cudnn: Whether to use CUDNN GPU optimised LSTM
        """

        self.name = self.__class__.__name__

        params = dict(raw_params)  # copy locally

        # Data parameters
        self.time_steps = int(params["total_time_steps"])
        self.input_size = int(params["input_size"])
        self.output_size = int(params["output_size"])
        self.category_counts = json.loads(str(params["category_counts"]))
        self.n_multiprocessing_workers = int(params["multiprocessing_workers"])

        # Relevant indices for TFT
        self._input_obs_loc = json.loads(str(params["input_obs_loc"]))
        self._static_input_loc = json.loads(str(params["static_input_loc"]))
        self._known_regular_input_idx = json.loads(str(params["known_regular_inputs"]))
        self._known_categorical_input_idx = json.loads(str(params["known_categorical_inputs"]))

        self.column_definition = params["column_definition"]

        # Network params
        self.quantiles = [0.1, 0.5, 0.9]
        self.use_cudnn = use_cudnn  # Whether to use GPU optimised LSTM
        self.hidden_layer_size = int(params["hidden_layer_size"])
        self.dropout_rate = float(params["dropout_rate"])
        self.max_gradient_norm = float(params["max_gradient_norm"])
        self.learning_rate = float(params["learning_rate"])
        self.minibatch_size = int(params["minibatch_size"])
        self.num_epochs = int(params["num_epochs"])
        self.early_stopping_patience = int(params["early_stopping_patience"])

        self.num_encoder_steps = int(params["num_encoder_steps"])
        self.num_stacks = int(params["stack_size"])
        self.num_heads = int(params["num_heads"])

        # Serialisation options
        self._temp_folder = os.path.join(params["model_folder"], "tmp")
        self.reset_temp_folder()

        # Extra components to store Tensorflow nodes for attention computations
        self._input_placeholder = None
        self._attention_components = None
        self._prediction_parts = None

        print("*** {} params ***".format(self.name))
        for k in params:
            print("# {} = {}".format(k, params[k]))

        # Build model
        self.model = self.build_model()

    def get_tft_embeddings(self, all_inputs):
        """Transforms raw inputs to embeddings.

        Applies linear transformation onto continuous variables and uses embeddings
        for categorical variables.

        Args:
          all_inputs: Inputs to transform

        Returns:
          Tensors for transformed inputs.
        """

        time_steps = self.time_steps

        # Sanity checks
        for i in self._known_regular_input_idx:
            if i in self._input_obs_loc:
                raise ValueError("Observation cannot be known a priori!")
        for i in self._input_obs_loc:
            if i in self._static_input_loc:
                raise ValueError("Observation cannot be static!")

        if all_inputs.get_shape().as_list()[-1] != self.input_size:
            raise ValueError(
                "Illegal number of inputs! Inputs observed={}, expected={}".format(
                    all_inputs.get_shape().as_list()[-1], self.input_size
                )
            )

        num_categorical_variables = len(self.category_counts)
        num_regular_variables = self.input_size - num_categorical_variables

        embedding_sizes = [self.hidden_layer_size for i, size in enumerate(self.category_counts)]

        embeddings = []
        for i in range(num_categorical_variables):

            embedding = tf.keras.Sequential(
                [
                    tf.keras.layers.InputLayer([time_steps]),
                    tf.keras.layers.Embedding(
                        self.category_counts[i], embedding_sizes[i], input_length=time_steps, dtype=tf.float32
                    ),
                ]
            )
            embeddings.append(embedding)

        regular_inputs, categorical_inputs = (
            all_inputs[:, :, :num_regular_variables],
            all_inputs[:, :, num_regular_variables:],
        )

        embedded_inputs = [embeddings[i](categorical_inputs[Ellipsis, i]) for i in range(num_categorical_variables)]

        # Static inputs
        if self._static_input_loc:
            static_inputs = [
                tf.keras.layers.Dense(self.hidden_layer_size)(regular_inputs[:, 0, i : i + 1])
                for i in range(num_regular_variables)
                if i in self._static_input_loc
            ] + [
                embedded_inputs[i][:, 0, :]
                for i in range(num_categorical_variables)
                if i + num_regular_variables in self._static_input_loc
            ]
            static_inputs = tf.keras.backend.stack(static_inputs, axis=1)

        else:
            static_inputs = None

        def convert_real_to_embedding(x):
            """Applies linear transformation for time-varying inputs."""
            return tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(self.hidden_layer_size))(x)

        # Targets
        obs_inputs = tf.keras.backend.stack(
            [convert_real_to_embedding(regular_inputs[Ellipsis, i : i + 1]) for i in self._input_obs_loc], axis=-1
        )

        # Observed (a prioir unknown) inputs
        wired_embeddings = []
        for i in range(num_categorical_variables):
            if i not in self._known_categorical_input_idx and i + num_regular_variables not in self._input_obs_loc:
                e = embeddings[i](categorical_inputs[:, :, i])
                wired_embeddings.append(e)

        unknown_inputs = []
        for i in range(regular_inputs.shape[-1]):
            if i not in self._known_regular_input_idx and i not in self._input_obs_loc:
                e = convert_real_to_embedding(regular_inputs[Ellipsis, i : i + 1])
                unknown_inputs.append(e)

        if unknown_inputs + wired_embeddings:
            unknown_inputs = tf.keras.backend.stack(unknown_inputs + wired_embeddings, axis=-1)
        else:
            unknown_inputs = None

        # A priori known inputs
        known_regular_inputs = [
            convert_real_to_embedding(regular_inputs[Ellipsis, i : i + 1])
            for i in self._known_regular_input_idx
            if i not in self._static_input_loc
        ]
        known_categorical_inputs = [
            embedded_inputs[i]
            for i in self._known_categorical_input_idx
            if i + num_regular_variables not in self._static_input_loc
        ]

        known_combined_layer = tf.keras.backend.stack(known_regular_inputs + known_categorical_inputs, axis=-1)

        return unknown_inputs, known_combined_layer, obs_inputs, static_inputs

    def _get_single_col_by_type(self, input_type):
        """Returns name of single column for input type."""

        return utils.get_single_col_by_input_type(input_type, self.column_definition)

    def training_data_cached(self):
        """Returns boolean indicating if training data has been cached."""

        return TFTDataCache.contains("train") and TFTDataCache.contains("valid")

    def cache_batched_data(self, data, cache_key, num_samples=-1):
        """Batches and caches data once for using during training.

        Args:
          data: Data to batch and cache
          cache_key: Key used for cache
          num_samples: Maximum number of samples to extract (-1 to use all data)
        """

        if num_samples > 0:
            TFTDataCache.update(self._batch_sampled_data(data, max_samples=num_samples), cache_key)
        else:
            TFTDataCache.update(self._batch_data(data), cache_key)

        print('Cached data "{}" updated'.format(cache_key))

    def _batch_sampled_data(self, data, max_samples):
        """Samples segments into a compatible format.

        Args:
          data: Sources data to sample and batch
          max_samples: Maximum number of samples in batch

        Returns:
          Dictionary of batched data with the maximum samples specified.
        """

        if max_samples < 1:
            raise ValueError("Illegal number of samples specified! samples={}".format(max_samples))

        id_col = self._get_single_col_by_type(InputTypes.ID)
        time_col = self._get_single_col_by_type(InputTypes.TIME)

        data.sort_values(by=[id_col, time_col], inplace=True)

        print("Getting valid sampling locations.")
        valid_sampling_locations = []
        split_data_map = {}
        for identifier, df in data.groupby(id_col):
            print("Getting locations for {}".format(identifier))
            num_entries = len(df)
            if num_entries >= self.time_steps:
                valid_sampling_locations += [
                    (identifier, self.time_steps + i) for i in range(num_entries - self.time_steps + 1)
                ]
            split_data_map[identifier] = df

        inputs = np.zeros((max_samples, self.time_steps, self.input_size))
        outputs = np.zeros((max_samples, self.time_steps, self.output_size))
        time = np.empty((max_samples, self.time_steps, 1), dtype=object)
        identifiers = np.empty((max_samples, self.time_steps, 1), dtype=object)

        if max_samples > 0 and len(valid_sampling_locations) > max_samples:
            print("Extracting {} samples...".format(max_samples))
            ranges = [
                valid_sampling_locations[i]
                for i in np.random.choice(len(valid_sampling_locations), max_samples, replace=False)
            ]
        else:
            print("Max samples={} exceeds # available segments={}".format(max_samples, len(valid_sampling_locations)))
            ranges = valid_sampling_locations

        id_col = self._get_single_col_by_type(InputTypes.ID)
        time_col = self._get_single_col_by_type(InputTypes.TIME)
        target_col = self._get_single_col_by_type(InputTypes.TARGET)
        input_cols = [tup[0] for tup in self.column_definition if tup[2] not in {InputTypes.ID, InputTypes.TIME}]

        for i, tup in enumerate(ranges):
            if (i + 1 % 1000) == 0:
                print(i + 1, "of", max_samples, "samples done...")
            identifier, start_idx = tup
            sliced = split_data_map[identifier].iloc[start_idx - self.time_steps : start_idx]
            inputs[i, :, :] = sliced[input_cols]
            outputs[i, :, :] = sliced[[target_col]]
            time[i, :, 0] = sliced[time_col]
            identifiers[i, :, 0] = sliced[id_col]

        sampled_data = {
            "inputs": inputs,
            "outputs": outputs[:, self.num_encoder_steps :, :],
            "active_entries": np.ones_like(outputs[:, self.num_encoder_steps :, :]),
            "time": time,
            "identifier": identifiers,
        }

        return sampled_data

    def _batch_data(self, data):
        """Batches data for training.

        Converts raw dataframe from a 2-D tabular format to a batched 3-D array
        to feed into Keras model.

        Args:
          data: DataFrame to batch

        Returns:
          Batched Numpy array with shape=(?, self.time_steps, self.input_size)
        """

        # Functions.
        def _batch_single_entity(input_data):
            time_steps = len(input_data)
            lags = self.time_steps
            x = input_data.values
            if time_steps >= lags:
                return np.stack([x[i : time_steps - (lags - 1) + i, :] for i in range(lags)], axis=1)

            else:
                return None

        id_col = self._get_single_col_by_type(InputTypes.ID)
        time_col = self._get_single_col_by_type(InputTypes.TIME)
        target_col = self._get_single_col_by_type(InputTypes.TARGET)
        input_cols = [tup[0] for tup in self.column_definition if tup[2] not in {InputTypes.ID, InputTypes.TIME}]

        data_map = {}
        for _, sliced in data.groupby(id_col):

            col_mappings = {"identifier": [id_col], "time": [time_col], "outputs": [target_col], "inputs": input_cols}

            for k in col_mappings:
                cols = col_mappings[k]
                arr = _batch_single_entity(sliced[cols].copy())

                if k not in data_map:
                    data_map[k] = [arr]
                else:
                    data_map[k].append(arr)

        # Combine all data
        for k in data_map:
            # Wendi: Avoid returning None when the length is not enough
            data_map[k] = np.concatenate([i for i in data_map[k] if i is not None], axis=0)

        # Shorten target so we only get decoder steps
        data_map["outputs"] = data_map["outputs"][:, self.num_encoder_steps :, :]

        active_entries = np.ones_like(data_map["outputs"])
        if "active_entries" not in data_map:
            data_map["active_entries"] = active_entries
        else:
            data_map["active_entries"].append(active_entries)

        return data_map

    def _get_active_locations(self, x):
        """Formats sample weights for Keras training."""
        return (np.sum(x, axis=-1) > 0.0) * 1.0

    def _build_base_graph(self):
        """Returns graph defining layers of the TFT."""

        # Size definitions.
        time_steps = self.time_steps
        combined_input_size = self.input_size
        encoder_steps = self.num_encoder_steps

        # Inputs.
        all_inputs = tf.keras.layers.Input(
            shape=(
                time_steps,
                combined_input_size,
            )
        )

        unknown_inputs, known_combined_layer, obs_inputs, static_inputs = self.get_tft_embeddings(all_inputs)

        # Isolate known and observed historical inputs.
        if unknown_inputs is not None:
            historical_inputs = concat(
                [
                    unknown_inputs[:, :encoder_steps, :],
                    known_combined_layer[:, :encoder_steps, :],
                    obs_inputs[:, :encoder_steps, :],
                ],
                axis=-1,
            )
        else:
            historical_inputs = concat(
                [known_combined_layer[:, :encoder_steps, :], obs_inputs[:, :encoder_steps, :]], axis=-1
            )

        # Isolate only known future inputs.
        future_inputs = known_combined_layer[:, encoder_steps:, :]

        def static_combine_and_mask(embedding):
            """Applies variable selection network to static inputs.

            Args:
              embedding: Transformed static inputs

            Returns:
              Tensor output for variable selection network
            """

            # Add temporal features
            _, num_static, _ = embedding.get_shape().as_list()

            flatten = tf.keras.layers.Flatten()(embedding)

            # Nonlinear transformation with gated residual network.
            mlp_outputs = gated_residual_network(
                flatten,
                self.hidden_layer_size,
                output_size=num_static,
                dropout_rate=self.dropout_rate,
                use_time_distributed=False,
                additional_context=None,
            )

            sparse_weights = tf.keras.layers.Activation("softmax")(mlp_outputs)
            sparse_weights = K.expand_dims(sparse_weights, axis=-1)

            trans_emb_list = []
            for i in range(num_static):
                e = gated_residual_network(
                    embedding[:, i : i + 1, :],
                    self.hidden_layer_size,
                    dropout_rate=self.dropout_rate,
                    use_time_distributed=False,
                )
                trans_emb_list.append(e)

            transformed_embedding = concat(trans_emb_list, axis=1)

            combined = tf.keras.layers.Multiply()([sparse_weights, transformed_embedding])

            static_vec = K.sum(combined, axis=1)

            return static_vec, sparse_weights

        static_encoder, static_weights = static_combine_and_mask(static_inputs)

        static_context_variable_selection = gated_residual_network(
            static_encoder, self.hidden_layer_size, dropout_rate=self.dropout_rate, use_time_distributed=False
        )
        static_context_enrichment = gated_residual_network(
            static_encoder, self.hidden_layer_size, dropout_rate=self.dropout_rate, use_time_distributed=False
        )
        static_context_state_h = gated_residual_network(
            static_encoder, self.hidden_layer_size, dropout_rate=self.dropout_rate, use_time_distributed=False
        )
        static_context_state_c = gated_residual_network(
            static_encoder, self.hidden_layer_size, dropout_rate=self.dropout_rate, use_time_distributed=False
        )

        def lstm_combine_and_mask(embedding):
            """Apply temporal variable selection networks.

            Args:
              embedding: Transformed inputs.

            Returns:
              Processed tensor outputs.
            """

            # Add temporal features
            _, time_steps, embedding_dim, num_inputs = embedding.get_shape().as_list()

            flatten = K.reshape(embedding, [-1, time_steps, embedding_dim * num_inputs])

            expanded_static_context = K.expand_dims(static_context_variable_selection, axis=1)

            # Variable selection weights
            mlp_outputs, static_gate = gated_residual_network(
                flatten,
                self.hidden_layer_size,
                output_size=num_inputs,
                dropout_rate=self.dropout_rate,
                use_time_distributed=True,
                additional_context=expanded_static_context,
                return_gate=True,
            )

            sparse_weights = tf.keras.layers.Activation("softmax")(mlp_outputs)
            sparse_weights = tf.expand_dims(sparse_weights, axis=2)

            # Non-linear Processing & weight application
            trans_emb_list = []
            for i in range(num_inputs):
                grn_output = gated_residual_network(
                    embedding[Ellipsis, i],
                    self.hidden_layer_size,
                    dropout_rate=self.dropout_rate,
                    use_time_distributed=True,
                )
                trans_emb_list.append(grn_output)

            transformed_embedding = stack(trans_emb_list, axis=-1)

            combined = tf.keras.layers.Multiply()([sparse_weights, transformed_embedding])
            temporal_ctx = K.sum(combined, axis=-1)

            return temporal_ctx, sparse_weights, static_gate

        historical_features, historical_flags, _ = lstm_combine_and_mask(historical_inputs)
        future_features, future_flags, _ = lstm_combine_and_mask(future_inputs)

        # LSTM layer
        def get_lstm(return_state):
            """Returns LSTM cell initialized with default parameters."""
            if self.use_cudnn:
                lstm = tf.keras.layers.CuDNNLSTM(
                    self.hidden_layer_size,
                    return_sequences=True,
                    return_state=return_state,
                    stateful=False,
                )
            else:
                lstm = tf.keras.layers.LSTM(
                    self.hidden_layer_size,
                    return_sequences=True,
                    return_state=return_state,
                    stateful=False,
                    # Additional params to ensure LSTM matches CuDNN, See TF 2.0 :
                    # (https://www.tensorflow.org/api_docs/python/tf/keras/layers/LSTM)
                    activation="tanh",
                    recurrent_activation="sigmoid",
                    recurrent_dropout=0,
                    unroll=False,
                    use_bias=True,
                )
            return lstm

        history_lstm, state_h, state_c = get_lstm(return_state=True)(
            historical_features, initial_state=[static_context_state_h, static_context_state_c]
        )

        future_lstm = get_lstm(return_state=False)(future_features, initial_state=[state_h, state_c])

        lstm_layer = concat([history_lstm, future_lstm], axis=1)

        # Apply gated skip connection
        input_embeddings = concat([historical_features, future_features], axis=1)

        lstm_layer, _ = apply_gating_layer(lstm_layer, self.hidden_layer_size, self.dropout_rate, activation=None)
        temporal_feature_layer = add_and_norm([lstm_layer, input_embeddings])

        # Static enrichment layers
        expanded_static_context = K.expand_dims(static_context_enrichment, axis=1)
        enriched, _ = gated_residual_network(
            temporal_feature_layer,
            self.hidden_layer_size,
            dropout_rate=self.dropout_rate,
            use_time_distributed=True,
            additional_context=expanded_static_context,
            return_gate=True,
        )

        # Decoder self attention
        self_attn_layer = InterpretableMultiHeadAttention(
            self.num_heads, self.hidden_layer_size, dropout=self.dropout_rate
        )

        mask = get_decoder_mask(enriched)
        x, self_att = self_attn_layer(enriched, enriched, enriched, mask=mask)

        x, _ = apply_gating_layer(x, self.hidden_layer_size, dropout_rate=self.dropout_rate, activation=None)
        x = add_and_norm([x, enriched])

        # Nonlinear processing on outputs
        decoder = gated_residual_network(
            x, self.hidden_layer_size, dropout_rate=self.dropout_rate, use_time_distributed=True
        )

        # Final skip connection
        decoder, _ = apply_gating_layer(decoder, self.hidden_layer_size, activation=None)
        transformer_layer = add_and_norm([decoder, temporal_feature_layer])

        # Attention components for explainability
        attention_components = {
            # Temporal attention weights
            "decoder_self_attn": self_att,
            # Static variable selection weights
            "static_flags": static_weights[Ellipsis, 0],
            # Variable selection weights of past inputs
            "historical_flags": historical_flags[Ellipsis, 0, :],
            # Variable selection weights of future inputs
            "future_flags": future_flags[Ellipsis, 0, :],
        }

        return transformer_layer, all_inputs, attention_components

    def build_model(self):
        """Build model and defines training losses.

        Returns:
          Fully defined Keras model.
        """

        with tf.variable_scope(self.name):

            transformer_layer, all_inputs, attention_components = self._build_base_graph()

            outputs = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(self.output_size * len(self.quantiles)))(
                transformer_layer[Ellipsis, self.num_encoder_steps :, :]
            )

            self._attention_components = attention_components

            adam = tf.keras.optimizers.Adam(lr=self.learning_rate, clipnorm=self.max_gradient_norm)

            model = tf.keras.Model(inputs=all_inputs, outputs=outputs)

            print(model.summary())

            valid_quantiles = self.quantiles
            output_size = self.output_size

            class QuantileLossCalculator:
                """Computes the combined quantile loss for prespecified quantiles.

                Attributes:
                  quantiles: Quantiles to compute losses
                """

                def __init__(self, quantiles):
                    """Initializes computer with quantiles for loss calculations.

                    Args:
                      quantiles: Quantiles to use for computations.
                    """
                    self.quantiles = quantiles

                def quantile_loss(self, a, b):
                    """Returns quantile loss for specified quantiles.

                    Args:
                      a: Targets
                      b: Predictions
                    """
                    quantiles_used = set(self.quantiles)

                    loss = 0.0
                    for i, quantile in enumerate(valid_quantiles):
                        if quantile in quantiles_used:
                            loss += utils.tensorflow_quantile_loss(
                                a[Ellipsis, output_size * i : output_size * (i + 1)],
                                b[Ellipsis, output_size * i : output_size * (i + 1)],
                                quantile,
                            )
                    return loss

            quantile_loss = QuantileLossCalculator(valid_quantiles).quantile_loss

            model.compile(loss=quantile_loss, optimizer=adam, sample_weight_mode="temporal")

            self._input_placeholder = all_inputs

        return model

    def fit(self, train_df=None, valid_df=None):
        """Fits deep neural network for given training and validation data.

        Args:
          train_df: DataFrame for training data
          valid_df: DataFrame for validation data
        """

        print("*** Fitting {} ***".format(self.name))

        # Add relevant callbacks
        callbacks = [
            tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=self.early_stopping_patience, min_delta=1e-4),
            tf.keras.callbacks.ModelCheckpoint(
                filepath=self.get_keras_saved_path(self._temp_folder),
                monitor="val_loss",
                save_best_only=True,
                save_weights_only=True,
            ),
            tf.keras.callbacks.TerminateOnNaN(),
        ]

        print("Getting batched_data")
        if train_df is None:
            print("Using cached training data")
            train_data = TFTDataCache.get("train")
        else:
            train_data = self._batch_data(train_df)

        if valid_df is None:
            print("Using cached validation data")
            valid_data = TFTDataCache.get("valid")
        else:
            valid_data = self._batch_data(valid_df)

        print("Using keras standard fit")

        def _unpack(data):
            return data["inputs"], data["outputs"], self._get_active_locations(data["active_entries"])

        # Unpack without sample weights
        data, labels, active_flags = _unpack(train_data)
        val_data, val_labels, val_flags = _unpack(valid_data)

        all_callbacks = callbacks

        self.model.fit(
            x=data,
            y=np.concatenate([labels, labels, labels], axis=-1),
            sample_weight=active_flags,
            epochs=self.num_epochs,
            batch_size=self.minibatch_size,
            validation_data=(val_data, np.concatenate([val_labels, val_labels, val_labels], axis=-1), val_flags),
            callbacks=all_callbacks,
            shuffle=True,
            use_multiprocessing=True,
            workers=self.n_multiprocessing_workers,
        )

        # Load best checkpoint again
        tmp_checkpont = self.get_keras_saved_path(self._temp_folder)
        if os.path.exists(tmp_checkpont):
            self.load(self._temp_folder, use_keras_loadings=True)

        else:
            print("Cannot load from {}, skipping ...".format(self._temp_folder))

    def evaluate(self, data=None, eval_metric="loss"):
        """Applies evaluation metric to the training data.

        Args:
          data: Dataframe for evaluation
          eval_metric: Evaluation metic to return, based on model definition.

        Returns:
          Computed evaluation loss.
        """

        if data is None:
            print("Using cached validation data")
            raw_data = TFTDataCache.get("valid")
        else:
            raw_data = self._batch_data(data)

        inputs = raw_data["inputs"]
        outputs = raw_data["outputs"]
        active_entries = self._get_active_locations(raw_data["active_entries"])

        metric_values = self.model.evaluate(
            x=inputs,
            y=np.concatenate([outputs, outputs, outputs], axis=-1),
            sample_weight=active_entries,
            workers=16,
            use_multiprocessing=True,
        )

        metrics = pd.Series(metric_values, self.model.metrics_names)

        return metrics[eval_metric]

    def predict(self, df, return_targets=False):
        """Computes predictions for a given input dataset.

        Args:
          df: Input dataframe
          return_targets: Whether to also return outputs aligned with predictions to
            faciliate evaluation

        Returns:
          Input dataframe or tuple of (input dataframe, algined output dataframe).
        """

        data = self._batch_data(df)

        inputs = data["inputs"]
        time = data["time"]
        identifier = data["identifier"]
        outputs = data["outputs"]

        combined = self.model.predict(inputs, workers=16, use_multiprocessing=True, batch_size=self.minibatch_size)

        # Format output_csv
        if self.output_size != 1:
            raise NotImplementedError("Current version only supports 1D targets!")

        def format_outputs(prediction):
            """Returns formatted dataframes for prediction."""

            flat_prediction = pd.DataFrame(
                prediction[:, :, 0], columns=["t+{}".format(i) for i in range(self.time_steps - self.num_encoder_steps)]
            )
            cols = list(flat_prediction.columns)
            flat_prediction["forecast_time"] = time[:, self.num_encoder_steps - 1, 0]
            flat_prediction["identifier"] = identifier[:, 0, 0]

            # Arrange in order
            return flat_prediction[["forecast_time", "identifier"] + cols]

        # Extract predictions for each quantile into different entries
        process_map = {
            "p{}".format(int(q * 100)): combined[Ellipsis, i * self.output_size : (i + 1) * self.output_size]
            for i, q in enumerate(self.quantiles)
        }

        if return_targets:
            # Add targets if relevant
            process_map["targets"] = outputs

        return {k: format_outputs(process_map[k]) for k in process_map}

    def get_attention(self, df):
        """Computes TFT attention weights for a given dataset.

        Args:
          df: Input dataframe

        Returns:
            Dictionary of numpy arrays for temporal attention weights and variable
              selection weights, along with their identifiers and time indices
        """

        data = self._batch_data(df)
        inputs = data["inputs"]
        identifiers = data["identifier"]
        time = data["time"]

        def get_batch_attention_weights(input_batch):
            """Returns weights for a given minibatch of data."""
            input_placeholder = self._input_placeholder
            attention_weights = {}
            for k in self._attention_components:
                attention_weight = tf.keras.backend.get_session().run(
                    self._attention_components[k], {input_placeholder: input_batch.astype(np.float32)}
                )
                attention_weights[k] = attention_weight
            return attention_weights

        # Compute number of batches
        batch_size = self.minibatch_size
        n = inputs.shape[0]
        num_batches = n // batch_size
        if n - (num_batches * batch_size) > 0:
            num_batches += 1

        # Split up inputs into batches
        batched_inputs = [inputs[i * batch_size : (i + 1) * batch_size, Ellipsis] for i in range(num_batches)]

        # Get attention weights, while avoiding large memory increases
        attention_by_batch = [get_batch_attention_weights(batch) for batch in batched_inputs]
        attention_weights = {}
        for k in self._attention_components:
            attention_weights[k] = []
            for batch_weights in attention_by_batch:
                attention_weights[k].append(batch_weights[k])

            if len(attention_weights[k][0].shape) == 4:
                tmp = np.concatenate(attention_weights[k], axis=1)
            else:
                tmp = np.concatenate(attention_weights[k], axis=0)

            del attention_weights[k]
            gc.collect()
            attention_weights[k] = tmp

        attention_weights["identifiers"] = identifiers[:, 0, 0]
        attention_weights["time"] = time[:, :, 0]

        return attention_weights

    # Serialisation.
    def reset_temp_folder(self):
        """Deletes and recreates folder with temporary Keras training outputs."""
        print("Resetting temp folder...")
        utils.create_folder_if_not_exist(self._temp_folder)
        shutil.rmtree(self._temp_folder)
        os.makedirs(self._temp_folder)

    def get_keras_saved_path(self, model_folder):
        """Returns path to keras checkpoint."""
        return os.path.join(model_folder, "{}.check".format(self.name))

    def save(self, model_folder):
        """Saves optimal TFT weights.

        Args:
          model_folder: Location to serialze model.
        """
        # Allows for direct serialisation of tensorflow variables to avoid spurious
        # issue with Keras that leads to different performance evaluation results
        # when model is reloaded (https://github.com/keras-team/keras/issues/4875).

        utils.save(tf.keras.backend.get_session(), model_folder, cp_name=self.name, scope=self.name)

    def load(self, model_folder, use_keras_loadings=False):
        """Loads TFT weights.

        Args:
          model_folder: Folder containing serialized models.
          use_keras_loadings: Whether to load from Keras checkpoint.

        Returns:

        """
        if use_keras_loadings:
            # Loads temporary Keras model saved during training.
            serialisation_path = self.get_keras_saved_path(model_folder)
            print("Loading model from {}".format(serialisation_path))
            self.model.load_weights(serialisation_path)
        else:
            # Loads tensorflow graph for optimal models.
            utils.load(tf.keras.backend.get_session(), model_folder, cp_name=self.name, scope=self.name)

    @classmethod
    def get_hyperparm_choices(cls):
        """Returns hyperparameter ranges for random search."""
        return {
            "dropout_rate": [0.1, 0.2, 0.3, 0.4, 0.5, 0.7, 0.9],
            "hidden_layer_size": [10, 20, 40, 80, 160, 240, 320],
            "minibatch_size": [64, 128, 256],
            "learning_rate": [1e-4, 1e-3, 1e-2],
            "max_gradient_norm": [0.01, 1.0, 100.0],
            "num_heads": [1, 4],
            "stack_size": [1],
        }