text: #
  # Copyright (c) 2019-2021, NVIDIA CORPORATION.
  #
  # 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.
  #
  
  ###############################################################################
  #                             How these tests work                            #
  ###############################################################################
  #
  # This test file contains some unit tests and an integration test.
  #
  # The units tests use the same parameters with cuML and the reference
  # implementation to compare strict parity of specific components.
  #
  # The integration tests compare that, when fitting and forecasting separately,
  # our implementation performs better or approximately as good as the reference
  # (it mostly serves to test that we don't have any regression)
  #
  # Note that there are significant differences between our implementation and
  # the reference, and perfect parity cannot be expected for integration tests.
  
  import pytest
  
  import numpy as np
  import os
  import warnings
  
  import pandas as pd
  from scipy.optimize.optimize import approx_fprime
  from sklearn.model_selection import train_test_split
  import statsmodels.api as sm
  
  import cudf
  import cuml.tsa.arima as arima
  from cuml.common.input_utils import input_to_host_array
  
  from cuml.test.utils import stress_param
  
  
  ###############################################################################
  #                                  Test data                                  #
  ###############################################################################
  
  class ARIMAData:
      """Contains a dataset name and associated metadata
      """
  
      def __init__(self, batch_size, n_obs, n_test, dataset,
                   tolerance_integration, n_exog=0, dataset_exog=None):
          self.batch_size = batch_size
          self.n_obs = n_obs
          self.n_test = n_test
          self.dataset = dataset
          self.tolerance_integration = tolerance_integration
  
          self.n_exog = n_exog
          self.dataset_exog = dataset_exog
  
          self.n_train = n_obs - n_test
  
  
  # ARIMA(1,0,1) with intercept
  test_101c = ARIMAData(
      batch_size=8,
      n_obs=15,
      n_test=2,
      dataset="long_term_arrivals_by_citizenship",
      tolerance_integration=0.01
  )
  
  # ARIMA(0,0,2) with intercept
  test_002c = ARIMAData(
      batch_size=7,
      n_obs=20,
      n_test=2,
      dataset="net_migrations_auckland_by_age",
      tolerance_integration=0.05
  )
  
  # ARIMA(0,1,0) with intercept
  test_010c = ARIMAData(
      batch_size=4,
      n_obs=17,
      n_test=2,
      dataset="cattle",
      tolerance_integration=0.01
  )
  
  # ARIMA(1,1,0)
  test_110 = ARIMAData(
      batch_size=1,
      n_obs=137,
      n_test=5,
      dataset="police_recorded_crime",
      tolerance_integration=0.01
  )
  
  # ARIMA(0,1,1) with intercept
  test_011c = ARIMAData(
      batch_size=16,
      n_obs=28,
      n_test=2,
      dataset="deaths_by_region",
      tolerance_integration=0.05
  )
  
  # ARIMA(0,1,1) with intercept (exogenous variables)
  test_011c_exog = ARIMAData(
      batch_size=16,
      n_obs=28,
      n_test=2,
      dataset="endog_deaths_by_region_exog",
      tolerance_integration=0.05,
      n_exog=2,
      dataset_exog="exog_deaths_by_region_exog"
  )
  
  # ARIMA(1,2,1) with intercept
  test_121c = ARIMAData(
      batch_size=2,
      n_obs=137,
      n_test=10,
      dataset="population_estimate",
      tolerance_integration=0.01
  )
  
  # ARIMA(1,1,1) with intercept (missing observations)
  test_111c_missing = ARIMAData(
      batch_size=2,
      n_obs=137,
      n_test=10,
      dataset="population_estimate_missing",
      tolerance_integration=0.01
  )
  
  # ARIMA(1,0,1)(1,1,1)_4
  test_101_111_4 = ARIMAData(
      batch_size=3,
      n_obs=101,
      n_test=10,
      dataset="alcohol",
      tolerance_integration=0.01
  )
  
  # ARIMA(5,1,0)
  test_510 = ARIMAData(
      batch_size=3,
      n_obs=101,
      n_test=10,
      dataset="alcohol",
      tolerance_integration=0.02
  )
  
  # ARIMA(1,1,1)(2,0,0)_4 with intercept
  test_111_200_4c = ARIMAData(
      batch_size=14,
      n_obs=123,
      n_test=10,
      dataset="hourly_earnings_by_industry",
      tolerance_integration=0.01
  )
  
  # ARIMA(1,1,1)(2,0,0)_4 with intercept (missing observations)
  test_111_200_4c_missing = ARIMAData(
      batch_size=14,
      n_obs=123,
      n_test=10,
      dataset="hourly_earnings_by_industry_missing",
      tolerance_integration=0.01
  )
  
  # ARIMA(1,1,1)(2,0,0)_4 with intercept
  # (missing observations and exogenous variables)
  test_111_200_4c_missing_exog = ARIMAData(
      batch_size=14,
      n_obs=123,
      n_test=10,
      dataset="endog_hourly_earnings_by_industry_missing_exog",
      tolerance_integration=0.01,
      n_exog=2,
      dataset_exog="exog_hourly_earnings_by_industry_missing_exog",
  )
  
  # ARIMA(1,1,2)(0,1,2)_4
  test_112_012_4 = ARIMAData(
      batch_size=2,
      n_obs=179,
      n_test=10,
      dataset="passenger_movements",
      tolerance_integration=0.001
  )
  
  # ARIMA(1,1,1)(1,1,1)_12
  test_111_111_12 = ARIMAData(
      batch_size=12,
      n_obs=279,
      n_test=20,
      dataset="guest_nights_by_region",
      tolerance_integration=0.001
  )
  
  # ARIMA(1,1,1)(1,1,1)_12 (missing observations)
  test_111_111_12_missing = ARIMAData(
      batch_size=12,
      n_obs=279,
      n_test=20,
      dataset="guest_nights_by_region_missing",
      tolerance_integration=0.03
  )
  
  # ARIMA(1,1,1)(1,1,1)_12 (missing obs, exogenous variables, intercept)
  test_111_111_12c_missing_exog = ARIMAData(
      batch_size=12,
      n_obs=279,
      n_test=20,
      dataset="endog_guest_nights_by_region_missing_exog",
      tolerance_integration=0.001,
      n_exog=2,
      dataset_exog="exog_guest_nights_by_region_missing_exog",
  )
  
  # Dictionary matching a test case to a tuple of model parameters
  # (a test case could be used with different models)
  # (p, d, q, P, D, Q, s, k) -> ARIMAData
  test_data = [
      # ((1, 0, 1, 0, 0, 0, 0, 1), test_101c),
      ((0, 0, 2, 0, 0, 0, 0, 1), test_002c),
      ((0, 1, 0, 0, 0, 0, 0, 1), test_010c),
      ((1, 1, 0, 0, 0, 0, 0, 0), test_110),
      ((0, 1, 1, 0, 0, 0, 0, 1), test_011c),
      ((0, 1, 1, 0, 0, 0, 0, 1), test_011c_exog),
      ((1, 2, 1, 0, 0, 0, 0, 1), test_121c),
      ((1, 1, 1, 0, 0, 0, 0, 1), test_111c_missing),
      ((1, 0, 1, 1, 1, 1, 4, 0), test_101_111_4),
      ((5, 1, 0, 0, 0, 0, 0, 0), test_510),
      ((1, 1, 1, 2, 0, 0, 4, 1), test_111_200_4c),
      ((1, 1, 1, 2, 0, 0, 4, 1), test_111_200_4c_missing),
      ((1, 1, 1, 2, 0, 0, 4, 1), test_111_200_4c_missing_exog),
      ((1, 1, 2, 0, 1, 2, 4, 0), test_112_012_4),
      stress_param((1, 1, 1, 1, 1, 1, 12, 0), test_111_111_12),
      stress_param((1, 1, 1, 1, 1, 1, 12, 0), test_111_111_12_missing),
      stress_param((1, 0, 1, 1, 1, 1, 12, 1), test_111_111_12c_missing_exog),
  ]
  
  # Dictionary for lazy-loading of datasets
  # (name, dtype) -> (pandas dataframe, cuDF dataframe)
  lazy_data = {}
  
  # Dictionary for lazy-evaluation of reference fits
  # (p, d, q, P, D, Q, s, k, name, dtype) -> SARIMAXResults
  lazy_ref_fit = {}
  
  
  def extract_order(tup):
      """Extract the order from a tuple of parameters"""
      p, d, q, P, D, Q, s, k = tup
      return (p, d, q), (P, D, Q, s), k
  
  
  data_path = os.path.join(os.path.dirname(
      os.path.abspath(__file__)), 'ts_datasets')
  
  
  def get_dataset(data, dtype):
      """Load a dataset with a given dtype or return a previously loaded dataset
      """
      key = (data.dataset, np.dtype(dtype).name)
      if key not in lazy_data:
          y = pd.read_csv(
              os.path.join(data_path, "{}.csv".format(data.dataset)),
              usecols=range(1, data.batch_size + 1), dtype=dtype)
          y_train, y_test = train_test_split(y, test_size=data.n_test,
                                             shuffle=False)
          y_train_cudf = cudf.from_pandas(y_train).fillna(np.nan)
          y_test_cudf = cudf.from_pandas(y_test)
          if data.dataset_exog is not None:
              exog = pd.read_csv(
                  os.path.join(data_path, "{}.csv".format(data.dataset_exog)),
                  usecols=range(1, data.n_exog * data.batch_size + 1),
                  dtype=dtype)
              exog_past, exog_fut = train_test_split(exog, test_size=data.n_test,
                                                     shuffle=False)
              exog_past_cudf = cudf.from_pandas(exog_past).fillna(np.nan)
              exog_fut_cudf = cudf.from_pandas(exog_fut)
          else:
              exog_past, exog_past_cudf, exog_fut, exog_fut_cudf = [None]*4
          lazy_data[key] = (y_train, y_train_cudf, y_test, y_test_cudf,
                            exog_past, exog_past_cudf, exog_fut, exog_fut_cudf)
      return lazy_data[key]
  
  
  def get_ref_fit(data, order, seasonal_order, intercept, dtype):
      """Compute a reference fit of a dataset with the given parameters and dtype
      or return a previously computed fit
      """
      y_train, _, _, _, exog_past, *_ = get_dataset(data, dtype)
      key = order + seasonal_order + \
          (intercept, data.dataset, np.dtype(dtype).name)
      batch_size = y_train.shape[1]
      if key not in lazy_ref_fit:
          ref_model = [
              sm.tsa.SARIMAX(endog=y_train[y_train.columns[i]],
                             exog=exog_past[exog_past.columns[
                                 data.n_exog*i:data.n_exog*(i+1)]]
                             if exog_past is not None else None,
                             order=order,
                             seasonal_order=seasonal_order,
                             trend='c' if intercept else 'n')
              for i in range(batch_size)]
          with warnings.catch_warnings():
              warnings.filterwarnings("ignore")
              lazy_ref_fit[key] = [model.fit(disp=0) for model in ref_model]
      return lazy_ref_fit[key]
  
  
  ###############################################################################
  #                              Utility functions                              #
  ###############################################################################
  
  def mase(y_train, y_test, y_fc, s):
      y_train_np = input_to_host_array(y_train).array
      y_test_np = input_to_host_array(y_test).array
      y_fc_np = input_to_host_array(y_fc).array
  
      diff = np.abs(y_train_np[s:] - y_train_np[:-s])
      scale = np.nanmean(diff, axis=0)
  
      error = np.abs(y_fc_np - y_test_np).mean(axis=0)
      return np.mean(error / scale)
  
  
  def fill_interpolation(df_in):
      np_arr = df_in.to_numpy()
      for ib in range(np_arr.shape[1]):
          n = len(np_arr)
          start, end = -1, 0
          while start < n - 1:
              if not np.isnan(np_arr[start+1, ib]):
                  start += 1
                  end = start + 1
              elif end < n and np.isnan(np_arr[end, ib]):
                  end += 1
              else:
                  if start == -1:
                      np_arr[:end, ib] = np_arr[end, ib]
                  elif end == n:
                      np_arr[start+1:, ib] = np_arr[start, ib]
                  else:
                      for j in range(start+1, end):
                          coef = (j - start) / (end - start)
                          np_arr[j, ib] = (
                              (1. - coef) * np_arr[start, ib]
                              + coef * np_arr[end, ib]
                          )
                  start = end
                  end = start + 1
      return pd.DataFrame(np_arr, columns=df_in.columns)
  
  
  ###############################################################################
  #                                    Tests                                    #
  ###############################################################################
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  def test_integration(key, data, dtype):
      """Full integration test: estimate, fit, forecast
      """
      order, seasonal_order, intercept = extract_order(key)
      s = max(1, seasonal_order[3])
  
      y_train, y_train_cudf, y_test, _, _, exog_past_cudf, exog_fut, \
          exog_fut_cudf = get_dataset(data, dtype)
  
      # Get fit reference model
      ref_fits = get_ref_fit(data, order, seasonal_order, intercept, dtype)
  
      # Create and fit cuML model
      cuml_model = arima.ARIMA(endog=y_train_cudf,
                               exog=exog_past_cudf,
                               order=order,
                               seasonal_order=seasonal_order,
                               fit_intercept=intercept,
                               output_type='numpy')
      cuml_model.fit()
  
      # Predict
      y_fc_cuml = cuml_model.forecast(data.n_test, exog=exog_fut)
      y_fc_ref = np.zeros((data.n_test, data.batch_size))
      for i in range(data.batch_size):
          y_fc_ref[:, i] = ref_fits[i].get_prediction(
              data.n_train, data.n_obs - 1,
              exog=None if data.n_exog == 0 else
              exog_fut[exog_fut.columns[data.n_exog*i:data.n_exog*(i+1)]]
          ).predicted_mean
  
      # Compare results: MASE must be better or within the tolerance margin
      mase_ref = mase(y_train, y_test, y_fc_ref, s)
      mase_cuml = mase(y_train, y_test, y_fc_cuml, s)
      assert mase_cuml < mase_ref * (1. + data.tolerance_integration)
  
  
  def _statsmodels_to_cuml(ref_fits, cuml_model, order, seasonal_order,
                           intercept, dtype):
      """Utility function to transfer the parameters from a statsmodels'
      SARIMAXResults object to a cuML ARIMA object.
  
      .. note:: be cautious with the intercept, it is not always equivalent
          in statsmodels and cuML models (it depends on the order).
  
      """
      nb = cuml_model.batch_size
      N = cuml_model.complexity
      x = np.zeros(nb * N, dtype=np.float64)
  
      for ib in range(nb):
          x[ib*N:(ib+1)*N] = ref_fits[ib].params[:N]
  
      cuml_model.unpack(x)
  
  
  def _predict_common(key, data, dtype, start, end, num_steps=None, level=None,
                      simple_differencing=True):
      """Utility function used by test_predict and test_forecast to avoid
      code duplication.
      """
      order, seasonal_order, intercept = extract_order(key)
  
      _, y_train_cudf, _, _, _, exog_cudf, exog_fut, exog_fut_cudf \
          = get_dataset(data, dtype)
  
      # Get fit reference model
      ref_fits = get_ref_fit(data, order, seasonal_order, intercept, dtype)
  
      # Create cuML model
      cuml_model = arima.ARIMA(endog=y_train_cudf,
                               exog=exog_cudf,
                               order=order,
                               seasonal_order=seasonal_order,
                               fit_intercept=intercept,
                               output_type='numpy',
                               simple_differencing=simple_differencing)
  
      # Feed the parameters to the cuML model
      _statsmodels_to_cuml(ref_fits, cuml_model, order, seasonal_order,
                           intercept, dtype)
  
      # Predict or forecast
      # Reference (statsmodels)
      ref_preds = np.zeros((end - start, data.batch_size))
      for i in range(data.batch_size):
          ref_preds[:, i] = ref_fits[i].get_prediction(
              start, end - 1,
              exog=(
                  None if data.n_exog == 0 or end <= data.n_train else
                  exog_fut[exog_fut.columns[data.n_exog*i:data.n_exog*(i+1)]])
          ).predicted_mean
      if level is not None:
          ref_lower = np.zeros((end - start, data.batch_size))
          ref_upper = np.zeros((end - start, data.batch_size))
          for i in range(data.batch_size):
              temp_pred = ref_fits[i].get_forecast(
                  num_steps,
                  exog=(
                      None if data.n_exog == 0 else exog_fut[
                          exog_fut.columns[data.n_exog*i:data.n_exog*(i+1)]])
              )
              ci = temp_pred.summary_frame(alpha=1-level)
              ref_lower[:, i] = ci["mean_ci_lower"].to_numpy()
              ref_upper[:, i] = ci["mean_ci_upper"].to_numpy()
      # cuML
      if num_steps is None:
          cuml_pred = cuml_model.predict(
              start, end,
              exog=None if data.n_exog == 0 or end <= data.n_train
              else exog_fut_cudf)
      elif level is not None:
          cuml_pred, cuml_lower, cuml_upper = \
              cuml_model.forecast(num_steps, level, exog=exog_fut_cudf)
      else:
          cuml_pred = cuml_model.forecast(num_steps, exog=exog_fut_cudf)
  
      # Compare results
      np.testing.assert_allclose(cuml_pred, ref_preds, rtol=0.002, atol=0.01)
      if level is not None:
          np.testing.assert_allclose(
              cuml_lower, ref_lower, rtol=0.005, atol=0.01)
          np.testing.assert_allclose(
              cuml_upper, ref_upper, rtol=0.005, atol=0.01)
  
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  @pytest.mark.parametrize('simple_differencing', [True, False])
  def test_predict_in(key, data, dtype, simple_differencing):
      """Test in-sample prediction against statsmodels (with the same values
      for the model parameters)
      """
      _predict_common(key, data, dtype, data.n_train // 2, data.n_obs,
                      simple_differencing=simple_differencing)
  
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  @pytest.mark.parametrize('simple_differencing', [True, False])
  def test_predict_inout(key, data, dtype, simple_differencing):
      """Test in- and ouf-of-sample prediction against statsmodels (with the
      same values for the model parameters)
      """
      _predict_common(key, data, dtype, data.n_train // 2, data.n_train,
                      simple_differencing=simple_differencing)
  
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  @pytest.mark.parametrize('simple_differencing', [True, False])
  def test_forecast(key, data, dtype, simple_differencing):
      """Test out-of-sample forecasting against statsmodels (with the same
      values for the model parameters)
      """
      _predict_common(key, data, dtype, data.n_train, data.n_obs, data.n_test,
                      simple_differencing=simple_differencing)
  
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  @pytest.mark.parametrize('level', [0.5, 0.95])
  def test_intervals(key, data, dtype, level):
      """Test forecast confidence intervals against statsmodels (with the same
      values for the model parameters)
      """
      _predict_common(key, data, dtype, data.n_train, data.n_obs, data.n_test,
                      level)
  
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  @pytest.mark.parametrize('simple_differencing', [True, False])
  def test_loglikelihood(key, data, dtype, simple_differencing):
      """Test loglikelihood against statsmodels (with the same values for the
      model parameters)
      """
      order, seasonal_order, intercept = extract_order(key)
  
      _, y_train_cudf, _, _, _, exog_past_cudf, *_ = get_dataset(data, dtype)
  
      # Get fit reference model
      ref_fits = get_ref_fit(data, order, seasonal_order, intercept, dtype)
  
      # Create cuML model
      cuml_model = arima.ARIMA(endog=y_train_cudf,
                               exog=exog_past_cudf,
                               order=order,
                               seasonal_order=seasonal_order,
                               fit_intercept=intercept,
                               simple_differencing=simple_differencing)
  
      # Feed the parameters to the cuML model
      _statsmodels_to_cuml(ref_fits, cuml_model, order, seasonal_order,
                           intercept, dtype)
  
      # Compute loglikelihood
      cuml_llf = cuml_model.llf
      ref_llf = np.array([ref_fit.llf for ref_fit in ref_fits])
  
      # Compare results
      np.testing.assert_allclose(cuml_llf, ref_llf, rtol=0.01, atol=0.01)
  
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  def test_gradient(key, data, dtype):
      """
      Test batched gradient implementation against scipy non-batched
      gradient.
  
      .. note:: it doesn't test that the loglikelihood is correct!
      """
      order, seasonal_order, intercept = extract_order(key)
      p, _, q = order
      P, _, Q, _ = seasonal_order
      h = 1e-8
  
      _, y_train_cudf, _, _, _, exog_past_cudf, *_ = get_dataset(data, dtype)
  
      # Create cuML model
      cuml_model = arima.ARIMA(endog=y_train_cudf,
                               exog=exog_past_cudf,
                               order=order,
                               seasonal_order=seasonal_order,
                               fit_intercept=intercept)
  
      N = cuml_model.complexity
  
      # Get an estimate of the parameters and pack them into a vector
      cuml_model._estimate_x0()
      x = cuml_model.pack()
  
      # Compute the batched loglikelihood gradient
      batched_grad = cuml_model._loglike_grad(x, h)
  
      # Iterate over the batch to compute a reference gradient
      scipy_grad = np.zeros(N * data.batch_size)
      for i in range(data.batch_size):
          # Create a model with only the current series
          model_i = arima.ARIMA(endog=y_train_cudf[y_train_cudf.columns[i]],
                                exog=None if exog_past_cudf is None
                                else exog_past_cudf[exog_past_cudf.columns[
                                    data.n_exog*i:data.n_exog*(i+1)]],
                                order=order,
                                seasonal_order=seasonal_order,
                                fit_intercept=intercept)
  
          def f(x):
              return model_i._loglike(x)
  
          scipy_grad[N * i: N * (i + 1)] = \
              approx_fprime(x[N * i: N * (i + 1)], f, h)
  
      # Compare
      np.testing.assert_allclose(batched_grad, scipy_grad, rtol=0.001, atol=0.01)
  
  
  @pytest.mark.parametrize('key, data', test_data)
  @pytest.mark.parametrize('dtype', [np.float64])
  def test_start_params(key, data, dtype):
      """Test starting parameters against statsmodels
      """
      order, seasonal_order, intercept = extract_order(key)
  
      y_train, y_train_cudf, _, _, exog_past, exog_past_cudf, *_ \
          = get_dataset(data, dtype)
  
      # fillna for reference to match cuML initial estimation strategy
      y_train_nona = fill_interpolation(y_train)
  
      # Convert to numpy to avoid misaligned indices
      if exog_past is not None:
          exog_past_np = exog_past.to_numpy()
  
      # Create models
      cuml_model = arima.ARIMA(endog=y_train_cudf,
                               exog=exog_past_cudf,
                               order=order,
                               seasonal_order=seasonal_order,
                               fit_intercept=intercept)
      ref_model = [sm.tsa.SARIMAX(endog=y_train_nona[y_train_nona.columns[i]],
                                  exog=exog_past_np[
                                      :, i*data.n_exog:(i+1)*data.n_exog]
                                  if data.n_exog else None,
                                  order=order,
                                  seasonal_order=seasonal_order,
                                  trend='c' if intercept else 'n')
                   for i in range(data.batch_size)]
  
      # Estimate reference starting parameters
      N = cuml_model.complexity
      nb = data.batch_size
      x_ref = np.zeros(N * nb, dtype=dtype)
      for ib in range(nb):
          with warnings.catch_warnings():
              warnings.filterwarnings("ignore")
              x_ref[ib*N:(ib+1)*N] = ref_model[ib].start_params[:N]
  
      # Estimate cuML starting parameters
      cuml_model._estimate_x0()
      x_cuml = cuml_model.pack()
  
      # Compare results
      np.testing.assert_allclose(x_cuml, x_ref, rtol=0.001, atol=0.01)
  