# Path: epi_forecast_stat_mech/evaluation/plot_constants.py
# def model_colors(k):
#   return {
#       'observed': {
#           'color': 'black',
#           'linestyle': 'solid'
#       },
#       'ground_truth': {
#           'color': 'black',
#           'linestyle': 'dashed'
#       },
#       'None_vc_Linear': {
#           'color': '#929292',
#           'linestyle': 'solid'
#       },
#       'None_vc_MLP': {
#           'color': '#00d0db',
#           'linestyle': 'dashed'
#       },
#       'None_Gaussian_Linear': {
#           'color': '#ff00ff',
#           'linestyle': 'solid'
#       },
#       'None_Gaussian_MLP': {
#           'color': '#b6b6b6',
#           'linestyle': 'dashed'
#       },
#       'Laplace_vc_Linear': {
#           'color': '#494949',
#           'linestyle': 'solid'
#       },
#       'Laplace_vc_MLP': {
#           'color': '#6d6d6d',
#           'linestyle': 'dashed'
#       },
#       'Laplace_Gaussian_Linear': {
#           'color': '#242424',
#           'linestyle': 'solid'
#       },
#       'Laplace_Gaussian_MLP': {
#           'color': '#101010',
#           'linestyle': 'dashed'
#       },
#       'rtlive': {
#           'color': '#55FFFF',
#           'linestyle': 'solid'
#       },
#       'sparse_classic': {
#           'color': '#DD0000',
#           'linestyle': ':'
#       },
#       'sparse_gaussian': {
#           'color': '#880000',
#           'linestyle': ':'
#       },
#       'iterative_mean': {
#           'color': '#00DD00',
#           'linestyle': ':'
#       },
#       'iterative_randomforest': {
#           'color': '#0000DD',
#           'linestyle': ':'
#       },
#   }.get(k, {
#       'color': '#39FF14',
#       'linestyle': ':'
#   })
# 
# Path: epi_forecast_stat_mech/evaluation/plot_constants.py
# def model_colors(k):


from epi_forecast_stat_mech.evaluation.plot_constants import model_colors
from epi_forecast_stat_mech.evaluation.plot_constants import model_types
from matplotlib import pyplot as plt
import numpy as np
  Returns:
    None
  """
  data_to_plot = data_inf.isel(location=location_to_plot)

  observed_color_params = model_colors('observed')
  max_observed_time = min(predictions.dropna('time').time)
  plt.plot(
      data_to_plot.coords['time'].sel(time=(data_inf.time < max_observed_time)),
      data_to_plot.sel(time=(data_inf.time < max_observed_time)),
      **observed_color_params,
      label='observed')
  return None


def plot_ground_truth_data(data_inf, predictions, location_to_plot):
  """Helper function to plot the ground truth data at a location.

  Args:
    data_inf: an xr.DataArray representing the *true* new_infections with
      dimensions of (location, time).
    predictions: an xr.DataArray representing predicted new_infections with
      dimensions of (location, time, sample, model). Where model is the
       Estimator used to generate the predicted new_infections.
    location_to_plot: The coordinate integer of the location that you want to
      plot.
  Returns:
    None
  """
  data_to_plot = data_inf.isel(location=location_to_plot)

# Path: epi_forecast_stat_mech/evaluation/plot_constants.py
# def model_colors(k):
#   return {
#       'observed': {
#           'color': 'black',
#           'linestyle': 'solid'
#       },
#       'ground_truth': {
#           'color': 'black',
#           'linestyle': 'dashed'
#       },
#       'None_vc_Linear': {
#           'color': '#929292',
#           'linestyle': 'solid'
#       },
#       'None_vc_MLP': {
#           'color': '#00d0db',
#           'linestyle': 'dashed'
#       },
#       'None_Gaussian_Linear': {
#           'color': '#ff00ff',
#           'linestyle': 'solid'
#       },
#       'None_Gaussian_MLP': {
#           'color': '#b6b6b6',
#           'linestyle': 'dashed'
#       },
#       'Laplace_vc_Linear': {
#           'color': '#494949',
#           'linestyle': 'solid'
#       },
#       'Laplace_vc_MLP': {
#           'color': '#6d6d6d',
#           'linestyle': 'dashed'
#       },
#       'Laplace_Gaussian_Linear': {
#           'color': '#242424',
#           'linestyle': 'solid'
#       },
#       'Laplace_Gaussian_MLP': {
#           'color': '#101010',
#           'linestyle': 'dashed'
#       },
#       'rtlive': {
#           'color': '#55FFFF',
#           'linestyle': 'solid'
#       },
#       'sparse_classic': {
#           'color': '#DD0000',
#           'linestyle': ':'
#       },
#       'sparse_gaussian': {
#           'color': '#880000',
#           'linestyle': ':'
#       },
#       'iterative_mean': {
#           'color': '#00DD00',
#           'linestyle': ':'
#       },
#       'iterative_randomforest': {
#           'color': '#0000DD',
#           'linestyle': ':'
#       },
#   }.get(k, {
#       'color': '#39FF14',
#       'linestyle': ':'
#   })
# 
# Path: epi_forecast_stat_mech/evaluation/plot_constants.py
# def model_colors(k):


# Path: epi_forecast_stat_mech/evaluation/plot_predictions.py
from epi_forecast_stat_mech.evaluation.plot_constants import model_colors
from epi_forecast_stat_mech.evaluation.plot_constants import model_types
from matplotlib import pyplot as plt
import numpy as np

  Returns:
    None
  """
  data_to_plot = data_inf.isel(location=location_to_plot)

  observed_color_params = model_colors('observed')
  max_observed_time = min(predictions.dropna('time').time)
  plt.plot(
      data_to_plot.coords['time'].sel(time=(data_inf.time < max_observed_time)),
      data_to_plot.sel(time=(data_inf.time < max_observed_time)),
      **observed_color_params,
      label='observed')
  return None


def plot_ground_truth_data(data_inf, predictions, location_to_plot):
  """Helper function to plot the ground truth data at a location.

  Args:
    data_inf: an xr.DataArray representing the *true* new_infections with
      dimensions of (location, time).
    predictions: an xr.DataArray representing predicted new_infections with
      dimensions of (location, time, sample, model). Where model is the
       Estimator used to generate the predicted new_infections.
    location_to_plot: The coordinate integer of the location that you want to
      plot.
  Returns:
    None
  """
  data_to_plot = data_inf.isel(location=location_to_plot)

  ground_truth_color_params = model_colors('ground_truth')