text: from collections import OrderedDict
  import numpy as np
  import torch
  import torch.nn as nn
  import pytorch_lightning as pl
  import logging
  from neuralprophet.utils import (
      season_config_to_model_dims,
      regressors_config_to_model_dims,
      events_config_to_model_dims,
  )
  
  log = logging.getLogger("NP.time_net")
  
  
  def new_param(dims):
      """Create and initialize a new torch Parameter.
  
      Args:
          dims (list or tuple): desired dimensions of parameter
  
      Returns:
          initialized Parameter
      """
      if len(dims) > 1:
          return nn.Parameter(nn.init.xavier_normal_(torch.randn(dims)), requires_grad=True)
      else:
          return nn.Parameter(torch.nn.init.xavier_normal_(torch.randn([1] + dims)).squeeze(0), requires_grad=True)
  
  
  class TimeNet(pl.LightningModule):
      """Linear time regression fun and some not so linear fun.
  
      A modular model that models classic time-series components
      - trend
      - seasonality
      - auto-regression (as AR-Net)
      - covariates (as AR-Net)
      - apriori regressors
      - events and holidays
      by using Neural Network components.
      The Auto-regression and covariate components can be configured as a deeper network (AR-Net).
      """
  
      def __init__(
          self,
          config_trend=None,
          config_season=None,
          config_covar=None,
          config_regressors=None,
          config_events=None,
          config_holidays=None,
          n_forecasts=1,
          n_lags=0,
          num_hidden_layers=0,
          d_hidden=None,
      ):
          """
          Args:
              config_trend (configure.Trend):
              config_season (configure.Season):
              config_covar (OrderedDict):
              config_regressors (OrderedDict): Configs of regressors with mode and index.
              config_events (OrderedDict):
              config_holidays (OrderedDict):
              n_forecasts (int): number of steps to forecast. Aka number of model outputs.
              n_lags (int): number of previous steps of time series used as input. Aka AR-order.
                  0 (default): no auto-regression
              num_hidden_layers (int): number of hidden layers (for AR-Net)
                  0 (default): no hidden layers, corresponds to classic Auto-Regression
              d_hidden (int): dimensionality of hidden layers  (for AR-Net). ignored if no hidden layers.
                  None (default): sets to n_lags + n_forecasts
          """
          super(TimeNet, self).__init__()
          # General
          self.n_forecasts = n_forecasts
  
          # Bias
          self.bias = new_param(dims=[1])
  
          # Metrics live
          self.metrics_live = {}
  
          # Trend
          self.config_trend = config_trend
          if self.config_trend.growth in ["linear", "discontinuous"]:
              self.segmentwise_trend = self.config_trend.trend_reg == 0
              self.trend_k0 = new_param(dims=[1])
              if self.config_trend.n_changepoints > 0:
                  if self.config_trend.changepoints is None:
                      # create equidistant changepoint times, including zero.
                      linear_t = np.arange(self.config_trend.n_changepoints + 1).astype(float)
                      linear_t = linear_t / (self.config_trend.n_changepoints + 1)
                      self.config_trend.changepoints = self.config_trend.changepoints_range * linear_t
                  else:
                      self.config_trend.changepoints = np.insert(self.config_trend.changepoints, 0, 0.0)
                  self.trend_changepoints_t = torch.tensor(
                      self.config_trend.changepoints, requires_grad=False, dtype=torch.float
                  )
                  self.trend_deltas = new_param(dims=[self.config_trend.n_changepoints + 1])  # including first segment
                  if self.config_trend.growth == "discontinuous":
                      self.trend_m = new_param(dims=[self.config_trend.n_changepoints + 1])  # including first segment
  
          # Seasonalities
          self.config_season = config_season
          self.season_dims = season_config_to_model_dims(self.config_season)
          if self.season_dims is not None:
              if self.config_season.mode == "multiplicative" and self.config_trend is None:
                  log.error("Multiplicative seasonality requires trend.")
                  raise ValueError
              if self.config_season.mode not in ["additive", "multiplicative"]:
                  log.error(
                      "Seasonality Mode {} not implemented. Defaulting to 'additive'.".format(self.config_season.mode)
                  )
                  self.config_season.mode = "additive"
              self.season_params = nn.ParameterDict(
                  {name: new_param(dims=[dim]) for name, dim in self.season_dims.items()}
              )
              # self.season_params_vec = torch.cat([self.season_params[name] for name in self.season_params.keys()])
  
          # Events
          self.config_events = config_events
          self.config_holidays = config_holidays
          self.events_dims = events_config_to_model_dims(self.config_events, self.config_holidays)
          if self.events_dims is not None:
              n_additive_event_params = 0
              n_multiplicative_event_params = 0
              for event, configs in self.events_dims.items():
                  if configs["mode"] not in ["additive", "multiplicative"]:
                      log.error("Event Mode {} not implemented. Defaulting to 'additive'.".format(configs["mode"]))
                      self.events_dims[event]["mode"] = "additive"
                  if configs["mode"] == "additive":
                      n_additive_event_params += len(configs["event_indices"])
                  elif configs["mode"] == "multiplicative":
                      if self.config_trend is None:
                          log.error("Multiplicative events require trend.")
                          raise ValueError
                      n_multiplicative_event_params += len(configs["event_indices"])
              self.event_params = nn.ParameterDict(
                  {
                      "additive": new_param(dims=[n_additive_event_params]),
                      "multiplicative": new_param(dims=[n_multiplicative_event_params]),
                  }
              )
          else:
              self.config_events = None
              self.config_holidays = None
  
              # Autoregression
          self.n_lags = n_lags
          self.num_hidden_layers = num_hidden_layers
          self.d_hidden = n_lags + n_forecasts if d_hidden is None else d_hidden
          if self.n_lags > 0:
              self.ar_net = nn.ModuleList()
              d_inputs = self.n_lags
              for i in range(self.num_hidden_layers):
                  self.ar_net.append(nn.Linear(d_inputs, self.d_hidden, bias=True))
                  d_inputs = self.d_hidden
              self.ar_net.append(nn.Linear(d_inputs, self.n_forecasts, bias=False))
              for lay in self.ar_net:
                  nn.init.kaiming_normal_(lay.weight, mode="fan_in")
  
          # Covariates
          self.config_covar = config_covar
          if self.config_covar is not None:
              assert self.n_lags > 0
              self.covar_nets = nn.ModuleDict({})
              for covar in self.config_covar.keys():
                  covar_net = nn.ModuleList()
                  d_inputs = self.n_lags
                  if self.config_covar[covar].as_scalar:
                      d_inputs = 1
                  for i in range(self.num_hidden_layers):
                      covar_net.append(nn.Linear(d_inputs, self.d_hidden, bias=True))
                      d_inputs = self.d_hidden
                  covar_net.append(nn.Linear(d_inputs, self.n_forecasts, bias=False))
                  for lay in covar_net:
                      nn.init.kaiming_normal_(lay.weight, mode="fan_in")
                  self.covar_nets[covar] = covar_net
  
          ## Regressors
          self.config_regressors = config_regressors
          self.regressors_dims = regressors_config_to_model_dims(config_regressors)
          if self.regressors_dims is not None:
              n_additive_regressor_params = 0
              n_multiplicative_regressor_params = 0
              for name, configs in self.regressors_dims.items():
                  if configs["mode"] not in ["additive", "multiplicative"]:
                      log.error("Regressors mode {} not implemented. Defaulting to 'additive'.".format(configs["mode"]))
                      self.regressors_dims[name]["mode"] = "additive"
                  if configs["mode"] == "additive":
                      n_additive_regressor_params += 1
                  elif configs["mode"] == "multiplicative":
                      if self.config_trend is None:
                          log.error("Multiplicative regressors require trend.")
                          raise ValueError
                      n_multiplicative_regressor_params += 1
  
              self.regressor_params = nn.ParameterDict(
                  {
                      "additive": new_param(dims=[n_additive_regressor_params]),
                      "multiplicative": new_param(dims=[n_multiplicative_regressor_params]),
                  }
              )
          else:
              self.config_regressors = None
  
      @property
      def get_trend_deltas(self):
          """trend deltas for regularization.
  
          update if trend is modelled differently"""
          if self.config_trend is None or self.config_trend.n_changepoints < 1:
              return None
          elif self.segmentwise_trend:
              return self.trend_deltas - torch.cat((self.trend_k0, self.trend_deltas[:-1]))
          else:
              return self.trend_deltas
  
      @property
      def ar_weights(self):
          """sets property auto-regression weights for regularization. Update if AR is modelled differently"""
          return self.ar_net[0].weight
  
      def get_covar_weights(self, name):
          """sets property auto-regression weights for regularization. Update if AR is modelled differently"""
          return self.covar_nets[name][0].weight
  
      def get_event_weights(self, name):
          """
          Retrieve the weights of event features given the name
  
          Args:
              name (string): Event name
  
          Returns:
              event_param_dict (OrderedDict): Dict of the weights of all offsets corresponding
              to a particular event.
          """
  
          event_dims = self.events_dims[name]
          mode = event_dims["mode"]
  
          if mode == "multiplicative":
              event_params = self.event_params["multiplicative"]
          else:
              assert mode == "additive"
              event_params = self.event_params["additive"]
  
          event_param_dict = OrderedDict({})
          for event_delim, indices in zip(event_dims["event_delim"], event_dims["event_indices"]):
              event_param_dict[event_delim] = event_params[indices]
          return event_param_dict
  
      def get_reg_weights(self, name):
          """
          Retrieve the weights of regressor features given the name
  
          Args:
              name (string): Regressor name
  
          Returns:
              weight (torch.tensor): Weight corresponding to the given regressor
          """
  
          regressor_dims = self.regressors_dims[name]
          mode = regressor_dims["mode"]
          index = regressor_dims["regressor_index"]
  
          if mode == "additive":
              regressor_params = self.regressor_params["additive"]
          if mode == "multiplicative":
              regressor_params = self.regressor_params["multiplicative"]
  
          return regressor_params[index]
  
      def _piecewise_linear_trend(self, t):
          """Piecewise linear trend, computed segmentwise or with deltas.
  
          Args:
              t (torch tensor, float): normalized time of
                  dimensions (batch, n_forecasts)
  
          Returns:
              Trend component, same dimensions as input t
          """
          past_next_changepoint = t.unsqueeze(2) >= torch.unsqueeze(self.trend_changepoints_t[1:], dim=0)
          segment_id = torch.sum(past_next_changepoint, dim=2)
          current_segment = nn.functional.one_hot(segment_id, num_classes=self.config_trend.n_changepoints + 1)
  
          k_t = torch.sum(current_segment * torch.unsqueeze(self.trend_deltas, dim=0), dim=2)
  
          if not self.segmentwise_trend:
              previous_deltas_t = torch.sum(past_next_changepoint * torch.unsqueeze(self.trend_deltas[:-1], dim=0), dim=2)
              k_t = k_t + previous_deltas_t
  
          if self.config_trend.growth != "discontinuous":
              if self.segmentwise_trend:
                  deltas = self.trend_deltas[:] - torch.cat((self.trend_k0, self.trend_deltas[0:-1]))
              else:
                  deltas = self.trend_deltas
              gammas = -self.trend_changepoints_t[1:] * deltas[1:]
              m_t = torch.sum(past_next_changepoint * gammas, dim=2)
              if not self.segmentwise_trend:
                  m_t = m_t.detach()
          else:
              m_t = torch.sum(current_segment * torch.unsqueeze(self.trend_m, dim=0), dim=2)
  
          return (self.trend_k0 + k_t) * t + m_t
  
      def trend(self, t):
          """Computes trend based on model configuration.
  
          Args:
              t (torch tensor float): normalized time
                  dimensions (batch, n_forecasts)
  
          Returns:
              Trend component, same dimensions as input t
  
          """
          if self.config_trend.growth == "off":
              trend = torch.zeros_like(t)
          elif int(self.config_trend.n_changepoints) == 0:
              trend = self.trend_k0 * t
          else:
              trend = self._piecewise_linear_trend(t)
          return self.bias + trend
  
      def seasonality(self, features, name):
          """Compute single seasonality component.
  
          Args:
              features (torch tensor, float): features related to seasonality component
                  dims: (batch, n_forecasts, n_features)
              name (str): name of seasonality. for attributiun to corresponding model weights.
  
          Returns:
              forecast component of dims (batch, n_forecasts)
          """
          return torch.sum(features * torch.unsqueeze(self.season_params[name], dim=0), dim=2)
  
      def all_seasonalities(self, s):
          """Compute all seasonality components.
  
          Args:
              s (dict(torch tensor, float)): dict of named seasonalities (keys) with their features (values)
                  dims of each dict value: (batch, n_forecasts, n_features)
  
          Returns:
              forecast component of dims (batch, n_forecasts)
          """
          x = torch.zeros(s[list(s.keys())[0]].shape[:2])
          for name, features in s.items():
              x = x + self.seasonality(features, name)
          return x
  
      def scalar_features_effects(self, features, params, indices=None):
          """
          Computes events component of the model
          Args:
              features (torch tensor, float): features (either additive or multiplicative) related to event component
                  dims: (batch, n_forecasts, n_features)
              params (nn.Parameter): params (either additive or multiplicative) related to events
              indices (list of int): indices in the feature tensors related to a particular event
          Returns:
              forecast component of dims (batch, n_forecasts)
          """
          if indices is not None:
              features = features[:, :, indices]
              params = params[indices]
  
          return torch.sum(features * torch.unsqueeze(params, dim=0), dim=2)
  
      def auto_regression(self, lags):
          """Computes auto-regessive model component AR-Net.
  
          Args:
              lags (torch tensor, float): previous times series values.
                  dims: (batch, n_lags)
  
          Returns:
              forecast component of dims: (batch, n_forecasts)
          """
          x = lags
          for i in range(self.num_hidden_layers + 1):
              if i > 0:
                  x = nn.functional.relu(x)
              x = self.ar_net[i](x)
          return x
  
      def covariate(self, lags, name):
          """Compute single covariate component.
  
          Args:
              lags (torch tensor, float): lagged values of covariate
                  dims: (batch, n_lags)
              name (str): name of covariate. for attributiun to corresponding model weights.
  
          Returns:
              forecast component of dims (batch, n_forecasts)
          """
          x = lags
          for i in range(self.num_hidden_layers + 1):
              if i > 0:
                  x = nn.functional.relu(x)
              x = self.covar_nets[name][i](x)
          return x
  
      def all_covariates(self, covariates):
          """Compute all covariate components.
  
          Args:
              covariates (dict(torch tensor, float)): dict of named covariates (keys) with their features (values)
                  dims of each dict value: (batch, n_lags)
  
          Returns:
              forecast component of dims (batch, n_forecasts)
          """
          for i, name in enumerate(covariates.keys()):
              if i == 0:
                  x = self.covariate(lags=covariates[name], name=name)
              if i > 0:
                  x = x + self.covariate(lags=covariates[name], name=name)
          return x
  
      def forward(self, inputs):
          """This method defines the model forward pass.
  
          Time input is required. Minimum model setup is a linear trend.
          Args:
              inputs (dict):
                  time (torch tensor float): normalized time
                      dims: (batch, n_forecasts)
                  lags (torch tensor, float): previous times series values.
                      dims: (batch, n_lags)
                  seasonalities (dict(torch tensor, float)): dict of named seasonalities (keys) with their features (values)
                      dims of each dict value: (batch, n_forecasts, n_features)
                  covariates (dict(torch tensor, float)): dict of named covariates (keys) with their features (values)
                      dims of each dict value: (batch, n_lags)
                  events (torch tensor, float): all event features
                      dims: (batch, n_forecasts, n_features)
                  regressors (torch tensor, float): all regressor features
                      dims: (batch, n_forecasts, n_features)
          Returns:
              forecast of dims (batch, n_forecasts)
          """
          additive_components = torch.zeros_like(inputs["time"])
          multiplicative_components = torch.zeros_like(inputs["time"])
  
          if "lags" in inputs:
              additive_components += self.auto_regression(lags=inputs["lags"])
          # else: assert self.n_lags == 0
  
          if "covariates" in inputs:
              additive_components += self.all_covariates(covariates=inputs["covariates"])
  
          if "seasonalities" in inputs:
              s = self.all_seasonalities(s=inputs["seasonalities"])
              if self.config_season.mode == "additive":
                  additive_components += s
              elif self.config_season.mode == "multiplicative":
                  multiplicative_components += s
  
          if "events" in inputs:
              if "additive" in inputs["events"].keys():
                  additive_components += self.scalar_features_effects(
                      inputs["events"]["additive"], self.event_params["additive"]
                  )
              if "multiplicative" in inputs["events"].keys():
                  multiplicative_components += self.scalar_features_effects(
                      inputs["events"]["multiplicative"], self.event_params["multiplicative"]
                  )
  
          if "regressors" in inputs:
              if "additive" in inputs["regressors"].keys():
                  additive_components += self.scalar_features_effects(
                      inputs["regressors"]["additive"], self.regressor_params["additive"]
                  )
              if "multiplicative" in inputs["regressors"].keys():
                  multiplicative_components += self.scalar_features_effects(
                      inputs["regressors"]["multiplicative"], self.regressor_params["multiplicative"]
                  )
  
          trend = self.trend(t=inputs["time"])
          out = trend + additive_components + trend.detach() * multiplicative_components
          return out
  
      def set_optimizer(self, optimizer):
          self.optimizer = optimizer
  
      def set_scheduler(self, scheduler):
          self.scheduler = scheduler
  
      def set_loss_func(self, loss_func):
          self.loss_func = loss_func
  
      def set_forecaster(self, self_forecaster):
          self.forecaster = self_forecaster
  
      def training_step(self, batch, batch_idx):
          x, y = batch
          y_hat = self(x)
          loss = self.loss_func(y_hat, y)
          e = self.current_epoch
  
          loss, reg_loss = self.forecaster._add_batch_regualarizations(
              loss, e, batch_idx / float(self.forecaster.loader_size)
          )
          self.log("train_loss", loss)
          self.forecaster.metrics.update(
              predicted=y_hat.detach(), target=y.detach(), values={"Loss": loss, "RegLoss": reg_loss}
          )
          return loss
  
      def validation_step(self, batch, batch_idx):
          x, y = batch
          y_hat = self(x)
          loss = self.loss_func(y_hat, y)
          self.log("val_loss", loss)
          self.forecaster.val_metrics.update(predicted=y_hat.detach(), target=y.detach())
  
          return loss
  
      def test_step(self, batch, batch_idx):
          x, y = batch
          y_hat = self(x)
          loss = self.loss_func(y_hat, y)
          self.log("test_loss", loss)
          self.forecaster.test_metrics.update(predicted=y_hat.detach(), target=y.detach())
  
          return loss
  
      def training_epoch_end(self, outputs):
          epoch_metrics = self.forecaster.metrics.compute(save=True)
          self.metrics_live["{}".format(list(epoch_metrics)[0])] = epoch_metrics[list(epoch_metrics)[0]]
  
          self.forecaster.metrics.reset()
  
      def validation_epoch_end(self, validation_step_outputs):
          val_epoch_metrics = self.forecaster.val_metrics.compute(save=True)
          self.metrics_live["val_{}".format(list(val_epoch_metrics)[0])] = val_epoch_metrics[list(val_epoch_metrics)[0]]
          self.forecaster.val_metrics.reset()
  
      def configure_optimizers(self):
          return [self.optimizer], [self.scheduler]
  
      def compute_components(self, inputs):
          """This method returns the values of each model component.
  
          Time input is required. Minimum model setup is a linear trend.
          Args:
              inputs (dict):
                  time (torch tensor float): normalized time
                      dims: (batch, n_forecasts)
                  lags (torch tensor, float): previous times series values.
                      dims: (batch, n_lags)
                  seasonalities (dict(torch tensor, float)): dict of named seasonalities (keys) with their features (values)
                      dims of each dict value: (batch, n_forecasts, n_features)
                  covariates (dict(torch tensor, float)): dict of named covariates (keys) with their features (values)
                      dims of each dict value: (batch, n_lags)
                  events (torch tensor, float): all event features
                      dims: (batch, n_forecasts, n_features)
          Returns:
              dict of forecast_component: value
                  with elements of dims (batch, n_forecasts)
          """
          components = {}
          components["trend"] = self.trend(t=inputs["time"])
          if self.config_trend is not None and "seasonalities" in inputs:
              for name, features in inputs["seasonalities"].items():
                  components["season_{}".format(name)] = self.seasonality(features=features, name=name)
          if self.n_lags > 0 and "lags" in inputs:
              components["ar"] = self.auto_regression(lags=inputs["lags"])
          if self.config_covar is not None and "covariates" in inputs:
              for name, lags in inputs["covariates"].items():
                  components["lagged_regressor_{}".format(name)] = self.covariate(lags=lags, name=name)
          if (self.config_events is not None or self.config_holidays is not None) and "events" in inputs:
              if "additive" in inputs["events"].keys():
                  components["events_additive"] = self.scalar_features_effects(
                      features=inputs["events"]["additive"], params=self.event_params["additive"]
                  )
              if "multiplicative" in inputs["events"].keys():
                  components["events_multiplicative"] = self.scalar_features_effects(
                      features=inputs["events"]["multiplicative"], params=self.event_params["multiplicative"]
                  )
              for event, configs in self.events_dims.items():
                  mode = configs["mode"]
                  indices = configs["event_indices"]
                  if mode == "additive":
                      features = inputs["events"]["additive"]
                      params = self.event_params["additive"]
                  else:
                      features = inputs["events"]["multiplicative"]
                      params = self.event_params["multiplicative"]
                  components["event_{}".format(event)] = self.scalar_features_effects(
                      features=features, params=params, indices=indices
                  )
          if self.config_regressors is not None and "regressors" in inputs:
              if "additive" in inputs["regressors"].keys():
                  components["future_regressors_additive"] = self.scalar_features_effects(
                      features=inputs["regressors"]["additive"], params=self.regressor_params["additive"]
                  )
              if "multiplicative" in inputs["regressors"].keys():
                  components["future_regressors_multiplicative"] = self.scalar_features_effects(
                      features=inputs["regressors"]["multiplicative"], params=self.regressor_params["multiplicative"]
                  )
              for regressor, configs in self.regressors_dims.items():
                  mode = configs["mode"]
                  index = []
                  index.append(configs["regressor_index"])
                  if mode == "additive":
                      features = inputs["regressors"]["additive"]
                      params = self.regressor_params["additive"]
                  else:
                      features = inputs["regressors"]["multiplicative"]
                      params = self.regressor_params["multiplicative"]
                  components["future_regressor_{}".format(regressor)] = self.scalar_features_effects(
                      features=features, params=params, indices=index
                  )
          return components
  
  
  class FlatNet(nn.Module):
      """
      Linear regression fun
      """
  
      def __init__(self, d_inputs, d_outputs):
          # Perform initialization of the pytorch superclass
          super(FlatNet, self).__init__()
          self.layers = nn.Sequential(
              nn.Linear(d_inputs, d_outputs),
          )
          nn.init.kaiming_normal_(self.layers[0].weight, mode="fan_in")
  
      def forward(self, x):
          return self.layers(x)
  
      @property
      def ar_weights(self):
          return self.model.layers[0].weight
  
  
  class DeepNet(nn.Module):
      """
      A simple, general purpose, fully connected network
      """
  
      def __init__(self, d_inputs, d_outputs, d_hidden=32, num_hidden_layers=0):
          # Perform initialization of the pytorch superclass
          super(DeepNet, self).__init__()
          self.layers = nn.ModuleList()
          for i in range(num_hidden_layers):
              self.layers.append(nn.Linear(d_inputs, d_hidden, bias=True))
              d_inputs = d_hidden
          self.layers.append(nn.Linear(d_inputs, d_outputs, bias=True))
          for lay in self.layers:
              nn.init.kaiming_normal_(lay.weight, mode="fan_in")
  
      def forward(self, x):
          """
          This method defines the network layering and activation functions
          """
          activation = nn.functional.relu
          for i in range(len(self.layers)):
              if i > 0:
                  x = activation(x)
              x = self.layers[i](x)
          return x
  
      @property
      def ar_weights(self):
          return self.layers[0].weight
  