# Path: epi_forecast_stat_mech/statistical_models/bic_model.py
# def logistic(x):
# def soft_nonzero(x, sharpness, threshold):
# def bic_adjustment(num_data_points, num_parameters):
#   def log_prior(self, parameters):
#   def _count_nonzero(self, params):
#   def log_likelihood(self, parameters, covariates, mechanistic_parameters):
#   def predict(self, parameters, covariates):
# class BICModel:
# 
# Path: epi_forecast_stat_mech/statistical_models/linear_model_test.py
# class LinearModelTest(parameterized.TestCase):
#   def assertBroadcastableTo(self, x, shape):
#   def get_model(self):
#   def testShapes(
#       self, covariate_shapes, mech_param_shapes, num_epidemics):


from absl.testing import parameterized
from epi_forecast_stat_mech.statistical_models import bic_model
from epi_forecast_stat_mech.statistical_models import linear_model_test
from absl.testing import absltest
import numpy as np

class BicLinearModelTest(linear_model_test.LinearModelTest):
  """We assert the BIC-wrapped model should pass the tests for `LinearModel`."""

  def get_model(self):
    base_model = super().get_model()
    return bic_model.BICModel(base_model=base_model)


class SoftNonzeroTest(parameterized.TestCase):

  @parameterized.parameters(
      dict(x=np.array([1, 2, 0, 0, 3], dtype=np.float32),
           sharpness=20,
           threshold=.1,
           nonzero_count=3),
      dict(x=np.array([[-1000, 200], [0, 0]], dtype=np.float32),
           sharpness=20,
           threshold=.1,
           nonzero_count=2),
      dict(x=np.array([0, 0, 0, 0], dtype=np.float32),
           sharpness=100,
           threshold=.1,
           nonzero_count=0),
      dict(x=np.arange(1000.),
           sharpness=20,
           threshold=.1,
           nonzero_count=999),
      )
  def testSoftNonzeroAccuracy(self, x, sharpness, threshold, nonzero_count):

# Path: epi_forecast_stat_mech/statistical_models/bic_model.py
# def logistic(x):
# def soft_nonzero(x, sharpness, threshold):
# def bic_adjustment(num_data_points, num_parameters):
#   def log_prior(self, parameters):
#   def _count_nonzero(self, params):
#   def log_likelihood(self, parameters, covariates, mechanistic_parameters):
#   def predict(self, parameters, covariates):
# class BICModel:
# 
# Path: epi_forecast_stat_mech/statistical_models/linear_model_test.py
# class LinearModelTest(parameterized.TestCase):
#   def assertBroadcastableTo(self, x, shape):
#   def get_model(self):
#   def testShapes(
#       self, covariate_shapes, mech_param_shapes, num_epidemics):


# Path: epi_forecast_stat_mech/statistical_models/bic_model_test.py
from absl.testing import parameterized
from epi_forecast_stat_mech.statistical_models import bic_model
from epi_forecast_stat_mech.statistical_models import linear_model_test
from absl.testing import absltest
import numpy as np


class BicLinearModelTest(linear_model_test.LinearModelTest):
  """We assert the BIC-wrapped model should pass the tests for `LinearModel`."""

  def get_model(self):
    base_model = super().get_model()
    return bic_model.BICModel(base_model=base_model)


class SoftNonzeroTest(parameterized.TestCase):

  @parameterized.parameters(
      dict(x=np.array([1, 2, 0, 0, 3], dtype=np.float32),
           sharpness=20,
           threshold=.1,
           nonzero_count=3),
      dict(x=np.array([[-1000, 200], [0, 0]], dtype=np.float32),
           sharpness=20,
           threshold=.1,
           nonzero_count=2),
      dict(x=np.array([0, 0, 0, 0], dtype=np.float32),
           sharpness=100,
           threshold=.1,
           nonzero_count=0),
      dict(x=np.arange(1000.),
           sharpness=20,
           threshold=.1,
           nonzero_count=999),
      )
  def testSoftNonzeroAccuracy(self, x, sharpness, threshold, nonzero_count):

    soft_nonzero_count = bic_model.soft_nonzero(x, sharpness, threshold).sum()