# 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
      observed_epidemic_record: initial epidemics trajectory to continue.
      dynamic_covariates: np array of time x dynamic_covariate.
      include_observed: whether to prepend the observed epidemics.

    Returns:
      trajectory of predicted infections of length `length` or
      `len(observed_epidemic_record) + length` if `include_observed` is True.
    """
    length = dynamic_covariates.shape[0] - len(observed_epidemic_record.t)
    assert length >= 0

    def _transition(state_and_rng, new_cases_and_params):
      state, rng = state_and_rng
      new_cases, params = new_cases_and_params
      if new_cases is not None:
        new_state = self._update_state(params, state, new_cases)
      else:
        current_rng, rng = jax.random.split(rng)
        intensity = self._intensity(params, state)
        new_cases = jnp.squeeze(
            self.new_infection_distribution(*intensity).sample(1, current_rng))
        new_state = self._update_state(params, state, new_cases)
      return (new_state, rng), new_cases

    observed_record_length = observed_epidemic_record.t.shape[0]
    time_series_length = observed_record_length + length
    time_dependent_params_shape = (time_series_length,
                                   len(self.step_based_encoded_param_names))
    time_dependent_params = jax.tree_multimap(jnp.broadcast_to, parameters,
                                              time_dependent_params_shape)

# 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

      observed_epidemic_record: initial epidemics trajectory to continue.
      dynamic_covariates: np array of time x dynamic_covariate.
      include_observed: whether to prepend the observed epidemics.

    Returns:
      trajectory of predicted infections of length `length` or
      `len(observed_epidemic_record) + length` if `include_observed` is True.
    """
    length = dynamic_covariates.shape[0] - len(observed_epidemic_record.t)
    assert length >= 0

    def _transition(state_and_rng, new_cases_and_params):
      state, rng = state_and_rng
      new_cases, params = new_cases_and_params
      if new_cases is not None:
        new_state = self._update_state(params, state, new_cases)
      else:
        current_rng, rng = jax.random.split(rng)
        intensity = self._intensity(params, state)
        new_cases = jnp.squeeze(
            self.new_infection_distribution(*intensity).sample(1, current_rng))
        new_state = self._update_state(params, state, new_cases)
      return (new_state, rng), new_cases

    observed_record_length = observed_epidemic_record.t.shape[0]
    time_series_length = observed_record_length + length
    time_dependent_params_shape = (time_series_length,
                                   len(self.step_based_encoded_param_names))
    time_dependent_params = jax.tree_multimap(jnp.broadcast_to, parameters,
                                              time_dependent_params_shape)

    start_params, params = utils.split_along_axis(time_dependent_params, 0, 1)