# Path: epi_forecast_stat_mech/data_model.py
# TIME_ZERO = np.datetime64('2019-12-31')
# def new_dataarray(parameters):
# def new_model(num_locations,
#               num_time_steps,
#               num_static_covariates,
#               num_dynamic_covariates=0):
# def _helper_shift_dataarray(shifts, array_to_shift):
# def _first_true(arr, axis):
# def shift_timeseries_by(data, shifts_all, enforce_positive_shifts: bool = True):
# def shift_timeseries(data, fraction_infected_limits, split_time):
# def validate_data(data,
#                   enable_regex_check=False,
#                   require_dynamics=False,
#                   require_samples=False,
#                   require_no_samples=False,
#                   require_integer_time=True,
#                   require_canonical_split_time=True):
# def validate_data_for_fit(data, **kwargs):
# def calculate_cumulative_infections(new_infections):
# def calculate_total_infections(new_infections):
# def calculate_infection_sums(ds):
# def _get_delta(series):
# def datetime_to_int(date, ds=None):
# def int_to_datetime(int_date, ds=None):
# def compute_integer_time(trajectories):
# def compute_numpy_index_time(trajectories):
# def convert_data_to_integer_time(trajectories, method='days_in_2020'):
# def set_sensible_canonical_split_time(data):


from absl.testing import absltest
from epi_forecast_stat_mech import data_model
import numpy as np
import xarray as xr

  def test_cumsum(self):
    new_da = data_model.calculate_cumulative_infections(self.ds.new_infections)
    np.testing.assert_array_equal(new_da.data,
                                  np.array([[1, 3, 6, 10], [0, 0, 0, 0]]))

  def test_sums(self):
    new_da = data_model.calculate_total_infections(self.ds.new_infections)
    np.testing.assert_array_equal(new_da.data, np.array([10, 0]))

  def test_all_sums(self):
    new_ds = data_model.calculate_infection_sums(self.ds)
    np.testing.assert_array_equal(new_ds.cumulative_infections.data,
                                  np.array([[1, 3, 6, 10], [0, 0, 0, 0]]))
    np.testing.assert_array_equal(new_ds.total_infections.data,
                                  np.array([10, 0]))


class TestShiftData(absltest.TestCase):

  def setUp(self):
    super(TestShiftData, self).setUp()
    self.num_locations = 3
    self.num_time_steps = 5
    self.dataarray = xr.DataArray(
        np.ones((self.num_locations, self.num_time_steps)),
        dims=['location', 'time'])
    self.shifts = 2 * np.ones(self.num_locations, dtype=int)

  def test_shift_array(self):

# Path: epi_forecast_stat_mech/data_model.py
# TIME_ZERO = np.datetime64('2019-12-31')
# def new_dataarray(parameters):
# def new_model(num_locations,
#               num_time_steps,
#               num_static_covariates,
#               num_dynamic_covariates=0):
# def _helper_shift_dataarray(shifts, array_to_shift):
# def _first_true(arr, axis):
# def shift_timeseries_by(data, shifts_all, enforce_positive_shifts: bool = True):
# def shift_timeseries(data, fraction_infected_limits, split_time):
# def validate_data(data,
#                   enable_regex_check=False,
#                   require_dynamics=False,
#                   require_samples=False,
#                   require_no_samples=False,
#                   require_integer_time=True,
#                   require_canonical_split_time=True):
# def validate_data_for_fit(data, **kwargs):
# def calculate_cumulative_infections(new_infections):
# def calculate_total_infections(new_infections):
# def calculate_infection_sums(ds):
# def _get_delta(series):
# def datetime_to_int(date, ds=None):
# def int_to_datetime(int_date, ds=None):
# def compute_integer_time(trajectories):
# def compute_numpy_index_time(trajectories):
# def convert_data_to_integer_time(trajectories, method='days_in_2020'):
# def set_sensible_canonical_split_time(data):


# Path: epi_forecast_stat_mech/tests/test_data_model.py
from absl.testing import absltest
from epi_forecast_stat_mech import data_model
import numpy as np
import xarray as xr


  def test_cumsum(self):
    new_da = data_model.calculate_cumulative_infections(self.ds.new_infections)
    np.testing.assert_array_equal(new_da.data,
                                  np.array([[1, 3, 6, 10], [0, 0, 0, 0]]))

  def test_sums(self):
    new_da = data_model.calculate_total_infections(self.ds.new_infections)
    np.testing.assert_array_equal(new_da.data, np.array([10, 0]))

  def test_all_sums(self):
    new_ds = data_model.calculate_infection_sums(self.ds)
    np.testing.assert_array_equal(new_ds.cumulative_infections.data,
                                  np.array([[1, 3, 6, 10], [0, 0, 0, 0]]))
    np.testing.assert_array_equal(new_ds.total_infections.data,
                                  np.array([10, 0]))


class TestShiftData(absltest.TestCase):

  def setUp(self):
    super(TestShiftData, self).setUp()
    self.num_locations = 3
    self.num_time_steps = 5
    self.dataarray = xr.DataArray(
        np.ones((self.num_locations, self.num_time_steps)),
        dims=['location', 'time'])
    self.shifts = 2 * np.ones(self.num_locations, dtype=int)

  def test_shift_array(self):

    shifted_data, _ = data_model._helper_shift_dataarray(