# Path: epi_forecast_stat_mech/utils.py
# def slice_along_axis(
#     inputs: Any,
#     axis: int,
#     idx: Union[slice, int]
# ):
# def split_along_axis(
#     inputs: Any,
#     axis: int,
#     split_idx: int,
#     post_squeeze_first: bool = True,
#     post_squeeze_second: bool = False
# ):
# 
# 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


import abc
import collections
import dataclasses
import flax.deprecated.nn.initializers
import jax
import jax.numpy as jnp
import jax.scipy.special
import numpy as np
import tensorflow_probability as tfp
from typing import Callable, Mapping, Text, Union
from epi_forecast_stat_mech import utils
from epi_forecast_stat_mech.statistical_models import probability as stat_prob
from flax.deprecated import nn
    cumulative_cases = observed_epidemics.cumulative_infections[-1]
    def _step(rng_and_cumulative, _):
      rng, cumulative_cases = rng_and_cumulative
      next_rng, rng = jax.random.split(rng)
      intensity = self.intensity(parameters, cumulative_cases)
      new_cases = jnp.squeeze(
          self.new_infection_distribution(intensity).sample(1, rng))
      cumulative_cases += new_cases
      return (next_rng, cumulative_cases), new_cases

    _, unroll = jax.lax.scan(_step, (rng, cumulative_cases), None, length)
    if include_observed:
      return jnp.concatenate([observed_epidemics.infections_over_time, unroll])
    return unroll


@dataclasses.dataclass
class GaussianModelPseudoLikelihood(GaussianModel):
  """Mechanistic model of Gaussian shape with a pseudo-likelihood criterion."""

  def __hash__(self):
    return id(self)

  def __eq__(self, other):
    return self is other

  def log_likelihood(self, parameters, epidemics):
    """Returns the (pseudo) log likelihood of `epidemics` given `parameters`."""
    intensity = self.intensity(parameters, epidemics.t)
    observations = epidemics.infections_over_time

# Path: epi_forecast_stat_mech/utils.py
# def slice_along_axis(
#     inputs: Any,
#     axis: int,
#     idx: Union[slice, int]
# ):
# def split_along_axis(
#     inputs: Any,
#     axis: int,
#     split_idx: int,
#     post_squeeze_first: bool = True,
#     post_squeeze_second: bool = False
# ):
# 
# 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/mechanistic_models/mechanistic_models.py
import abc
import collections
import dataclasses
import flax.deprecated.nn.initializers
import jax
import jax.numpy as jnp
import jax.scipy.special
import numpy as np
import tensorflow_probability as tfp
from typing import Callable, Mapping, Text, Union
from epi_forecast_stat_mech import utils
from epi_forecast_stat_mech.statistical_models import probability as stat_prob
from flax.deprecated import nn

    cumulative_cases = observed_epidemics.cumulative_infections[-1]
    def _step(rng_and_cumulative, _):
      rng, cumulative_cases = rng_and_cumulative
      next_rng, rng = jax.random.split(rng)
      intensity = self.intensity(parameters, cumulative_cases)
      new_cases = jnp.squeeze(
          self.new_infection_distribution(intensity).sample(1, rng))
      cumulative_cases += new_cases
      return (next_rng, cumulative_cases), new_cases

    _, unroll = jax.lax.scan(_step, (rng, cumulative_cases), None, length)
    if include_observed:
      return jnp.concatenate([observed_epidemics.infections_over_time, unroll])
    return unroll


@dataclasses.dataclass
class GaussianModelPseudoLikelihood(GaussianModel):
  """Mechanistic model of Gaussian shape with a pseudo-likelihood criterion."""

  def __hash__(self):
    return id(self)

  def __eq__(self, other):
    return self is other

  def log_likelihood(self, parameters, epidemics):
    """Returns the (pseudo) log likelihood of `epidemics` given `parameters`."""
    intensity = self.intensity(parameters, epidemics.t)
    observations = epidemics.infections_over_time

    log_probs = stat_prob.pseudo_poisson_log_probs(intensity, observations)