# Path: epi_forecast_stat_mech/constants.py
# def coord_units(key):
#   return {
#       'sample': 'arbitrary int',
#       'location': 'arbitrary int',
#       'time': 'simulation step',
#       'static_covariate': 'arbitrary int',
#   }.get(key)
# 
# Path: epi_forecast_stat_mech/constants.py
# def coordinates(key, number):
#   return {
#       'sample': np.arange(number),
#       'location': np.arange(number),
#       'time': np.arange(number),
#       'static_covariate': np.arange(number),
#   }.get(key)


from .constants import coord_units
from .constants import coordinates
from absl import logging
import numpy as np
import re
import xarray as xr
valid_dimensions = [
    'location', 'time', 'static_covariate', 'dynamic_covariate', 'model'
]

TIME_ZERO = np.datetime64('2019-12-31')


def new_dataarray(parameters):
  """Helper function to create xr.DataArrays.

  Create xr.DataArrays with consistent dimensions, coordinates, and units.

  Args:
    parameters: a dictionary of (dimension: number) pairs. Dimension must be in
      valid_dimensions. Number is the length of that axis.

  Returns:
    da: a xr.DataArray containing zeros along all specified dimensions
  """
  # Check all dimensions are valid
  for dim in parameters.keys():
    if dim not in valid_dimensions:
      raise ValueError('Unsupported parameter \'%s\'' % (dim))

  array_shape = tuple(parameters.values())
  dimensions = tuple(parameters.keys())
  coords = {}
  for dim in dimensions:
    # can change to something more sophisticated later
    if np.any(coordinates(dim, parameters[dim])):

# Path: epi_forecast_stat_mech/constants.py
# def coord_units(key):
#   return {
#       'sample': 'arbitrary int',
#       'location': 'arbitrary int',
#       'time': 'simulation step',
#       'static_covariate': 'arbitrary int',
#   }.get(key)
# 
# Path: epi_forecast_stat_mech/constants.py
# def coordinates(key, number):
#   return {
#       'sample': np.arange(number),
#       'location': np.arange(number),
#       'time': np.arange(number),
#       'static_covariate': np.arange(number),
#   }.get(key)


# Path: epi_forecast_stat_mech/data_model.py
from .constants import coord_units
from .constants import coordinates
from absl import logging
import numpy as np
import re
import xarray as xr

valid_dimensions = [
    'location', 'time', 'static_covariate', 'dynamic_covariate', 'model'
]

TIME_ZERO = np.datetime64('2019-12-31')


def new_dataarray(parameters):
  """Helper function to create xr.DataArrays.

  Create xr.DataArrays with consistent dimensions, coordinates, and units.

  Args:
    parameters: a dictionary of (dimension: number) pairs. Dimension must be in
      valid_dimensions. Number is the length of that axis.

  Returns:
    da: a xr.DataArray containing zeros along all specified dimensions
  """
  # Check all dimensions are valid
  for dim in parameters.keys():
    if dim not in valid_dimensions:
      raise ValueError('Unsupported parameter \'%s\'' % (dim))

  array_shape = tuple(parameters.values())
  dimensions = tuple(parameters.keys())
  coords = {}
  for dim in dimensions:
    # can change to something more sophisticated later
    if np.any(coordinates(dim, parameters[dim])):

      coords[dim] = coordinates(dim, parameters[dim])