# Path: epi_forecast_stat_mech/statistical_models/probability.py
# def softplus(x, floor, softness=1):
# def softplus_scale_estimate(
#     x, axis, min_scale, softness=1):
# def gaussian_with_softplus_scale_estimate(
#     x, axis, min_scale, mean=0, softness=1):
# def pseudo_poisson_log_probs(intensity, observations, axis=0):
# def soft_laplace_log_prob(x, scale=1):
# def laplace_prior(parameters, scale_parameter=1.):
# def log_mixed_laplace(x, a=0., b=0.5, include_constants=False):
# def soft_abs(x, delta):
# def log_soft_mixed_laplace(x, a=0., b=0.5, delta=0.01, include_constants=False):
# def log_soft_mixed_laplace_single(x,
#                                   a=0.,
#                                   b=0.5,
#                                   delta=0.01,
#                                   include_constants=False):
# def apply_to_recursive_dict_by_last_key(last_key_mapper, d, default):
# def _rec_apply_to_recursive_dict_by_last_key(last_key_mapper, d, last_key,
#                                              default):
# def log_soft_mixed_laplace_on_kernels(flax_dict):
# def bounded_derivative_log(x, baseline_constant):
# def gaussian_error_logprob_with_bottom_scale(error, bottom_scale, axis=0):
# LOG_SOFT_LAPLACE_INTEGRAL = 1.185495232349193
# 
# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


import collections
import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow_probability
from typing import Optional, Union
from epi_forecast_stat_mech.statistical_models import probability
from epi_forecast_stat_mech.statistical_models import tree_util
  def init(cls, covariate_dim, prediction_dim, intercept=None):
    # TODO(jamieas): initialize more intelligently.
    if intercept is None:
      intercept = jnp.zeros([1, prediction_dim])
    return cls(alpha=jnp.zeros([covariate_dim, prediction_dim]),
               intercept=intercept)


@dataclasses.dataclass
class LinearModel:
  """A linear model for predicting mechanistic parameters."""

  error_scale: Array = 1.
  # The scale of the Laplace prior placed on the parameters. A larger value
  # corresponds to a less informative prior. If `None`, then we use an improper
  # flat prior (effectively scale=infinity).
  # TODO(jamieas): make it easier to pass per-parameter scales so that we can
  # construct a data-dependent prior.
  laplace_prior_scale: Optional[Array] = None

  def log_prior(self, parameters):
    """Returns the log prior probability of `parameters`."""
    if self.laplace_prior_scale is None:
      return 0.
    return tree_util.tree_map(probability.soft_laplace_log_prob, parameters)

  def log_likelihood(self, parameters, covariates, mechanistic_parameters):
    """Returns the log likelihood of `mechanistic_parameters`."""
    predicted_mech_params = self.predict(parameters, covariates)
    gaussian = tfd.Normal(predicted_mech_params, self.error_scale)

# Path: epi_forecast_stat_mech/statistical_models/probability.py
# def softplus(x, floor, softness=1):
# def softplus_scale_estimate(
#     x, axis, min_scale, softness=1):
# def gaussian_with_softplus_scale_estimate(
#     x, axis, min_scale, mean=0, softness=1):
# def pseudo_poisson_log_probs(intensity, observations, axis=0):
# def soft_laplace_log_prob(x, scale=1):
# def laplace_prior(parameters, scale_parameter=1.):
# def log_mixed_laplace(x, a=0., b=0.5, include_constants=False):
# def soft_abs(x, delta):
# def log_soft_mixed_laplace(x, a=0., b=0.5, delta=0.01, include_constants=False):
# def log_soft_mixed_laplace_single(x,
#                                   a=0.,
#                                   b=0.5,
#                                   delta=0.01,
#                                   include_constants=False):
# def apply_to_recursive_dict_by_last_key(last_key_mapper, d, default):
# def _rec_apply_to_recursive_dict_by_last_key(last_key_mapper, d, last_key,
#                                              default):
# def log_soft_mixed_laplace_on_kernels(flax_dict):
# def bounded_derivative_log(x, baseline_constant):
# def gaussian_error_logprob_with_bottom_scale(error, bottom_scale, axis=0):
# LOG_SOFT_LAPLACE_INTEGRAL = 1.185495232349193
# 
# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


# Path: epi_forecast_stat_mech/statistical_models/linear_model.py
import collections
import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow_probability
from typing import Optional, Union
from epi_forecast_stat_mech.statistical_models import probability
from epi_forecast_stat_mech.statistical_models import tree_util

  def init(cls, covariate_dim, prediction_dim, intercept=None):
    # TODO(jamieas): initialize more intelligently.
    if intercept is None:
      intercept = jnp.zeros([1, prediction_dim])
    return cls(alpha=jnp.zeros([covariate_dim, prediction_dim]),
               intercept=intercept)


@dataclasses.dataclass
class LinearModel:
  """A linear model for predicting mechanistic parameters."""

  error_scale: Array = 1.
  # The scale of the Laplace prior placed on the parameters. A larger value
  # corresponds to a less informative prior. If `None`, then we use an improper
  # flat prior (effectively scale=infinity).
  # TODO(jamieas): make it easier to pass per-parameter scales so that we can
  # construct a data-dependent prior.
  laplace_prior_scale: Optional[Array] = None

  def log_prior(self, parameters):
    """Returns the log prior probability of `parameters`."""
    if self.laplace_prior_scale is None:
      return 0.
    return tree_util.tree_map(probability.soft_laplace_log_prob, parameters)

  def log_likelihood(self, parameters, covariates, mechanistic_parameters):
    """Returns the log likelihood of `mechanistic_parameters`."""
    predicted_mech_params = self.predict(parameters, covariates)
    gaussian = tfd.Normal(predicted_mech_params, self.error_scale)

    mech_params_array, _ = tree_util.pack(mechanistic_parameters)