# Path: pyhawkes/models.py
# class DiscreteTimeNetworkHawkesModelSpikeAndSlab(_DiscreteTimeNetworkHawkesModelBase, ModelGibbsSampling):
#     _weight_class           = SpikeAndSlabGammaWeights
#     _default_weight_hypers  = {}
# 
#     _network_class          = ErdosRenyiFixedSparsity
#     _default_network_hypers = {'p': 0.5,
#                                'allow_self_connections': True,
#                                'kappa': 1.0,
#                                'v': None, 'alpha': None, 'beta': None,}
# 
#     def resample_model(self):
#         """
#         Perform one iteration of the Gibbs sampling algorithm.
#         :return:
#         """
#         # Update the parents.
#         # THIS MUST BE DONE IMMEDIATELY FOLLOWING WEIGHT UPDATES!
#         for p in self.data_list:
#             p.resample()
# 
#         # Update the bias model given the parents assigned to the background
#         self.bias_model.resample(self.data_list)
# 
#         # Update the impulse model given the parents assignments
#         self.impulse_model.resample(self.data_list)
# 
#         # Update the network model
#         self.network.resample(data=(self.weight_model.A, self.weight_model.W))
# 
#         # Update the weight model given the parents assignments
#         self.weight_model.resample(self.data_list)
# 
#     def initialize_with_standard_model(self, standard_model):
#         super(DiscreteTimeNetworkHawkesModelSpikeAndSlab, self).\
#             initialize_with_standard_model(standard_model)
# 
#         # Update the parents.
#         for d in self.data_list:
#             d.resample()


import numpy as np
from pyhawkes.models import DiscreteTimeNetworkHawkesModelSpikeAndSlab

def test_compute_rate():
    K = 1
    T = 100
    dt = 1.0
    network_hypers = {'c': np.zeros(K, dtype=np.int), 'p': 1.0, 'kappa': 10.0, 'v': 10*5.0}
    true_model = DiscreteTimeNetworkHawkesModelSpikeAndSlab(K=K, dt=dt,
                                                            network_hypers=network_hypers)
    S,R = true_model.generate(T=T)

    print("Expected number of events: ", np.trapz(R, dt * np.arange(T), axis=0))
    print("Actual number of events:   ", S.sum(axis=0))

    print("Lambda0:  ", true_model.bias_model.lambda0)
    print("W:        ", true_model.weight_model.W)
    print("")

    R_test = true_model.compute_rate()
    assert np.allclose(R, R_test)

def test_generate_statistics():
    K = 1
    T = 100
    dt = 1.0
    network_hypers = {'c': np.zeros(K, dtype=np.int), 'p': 1.0, 'kappa': 10.0, 'v': 10*5.0}

# Path: pyhawkes/models.py
# class DiscreteTimeNetworkHawkesModelSpikeAndSlab(_DiscreteTimeNetworkHawkesModelBase, ModelGibbsSampling):
#     _weight_class           = SpikeAndSlabGammaWeights
#     _default_weight_hypers  = {}
# 
#     _network_class          = ErdosRenyiFixedSparsity
#     _default_network_hypers = {'p': 0.5,
#                                'allow_self_connections': True,
#                                'kappa': 1.0,
#                                'v': None, 'alpha': None, 'beta': None,}
# 
#     def resample_model(self):
#         """
#         Perform one iteration of the Gibbs sampling algorithm.
#         :return:
#         """
#         # Update the parents.
#         # THIS MUST BE DONE IMMEDIATELY FOLLOWING WEIGHT UPDATES!
#         for p in self.data_list:
#             p.resample()
# 
#         # Update the bias model given the parents assigned to the background
#         self.bias_model.resample(self.data_list)
# 
#         # Update the impulse model given the parents assignments
#         self.impulse_model.resample(self.data_list)
# 
#         # Update the network model
#         self.network.resample(data=(self.weight_model.A, self.weight_model.W))
# 
#         # Update the weight model given the parents assignments
#         self.weight_model.resample(self.data_list)
# 
#     def initialize_with_standard_model(self, standard_model):
#         super(DiscreteTimeNetworkHawkesModelSpikeAndSlab, self).\
#             initialize_with_standard_model(standard_model)
# 
#         # Update the parents.
#         for d in self.data_list:
#             d.resample()


# Path: test/test_model.py
import numpy as np
from pyhawkes.models import DiscreteTimeNetworkHawkesModelSpikeAndSlab


def test_compute_rate():
    K = 1
    T = 100
    dt = 1.0
    network_hypers = {'c': np.zeros(K, dtype=np.int), 'p': 1.0, 'kappa': 10.0, 'v': 10*5.0}
    true_model = DiscreteTimeNetworkHawkesModelSpikeAndSlab(K=K, dt=dt,
                                                            network_hypers=network_hypers)
    S,R = true_model.generate(T=T)

    print("Expected number of events: ", np.trapz(R, dt * np.arange(T), axis=0))
    print("Actual number of events:   ", S.sum(axis=0))

    print("Lambda0:  ", true_model.bias_model.lambda0)
    print("W:        ", true_model.weight_model.W)
    print("")

    R_test = true_model.compute_rate()
    assert np.allclose(R, R_test)

def test_generate_statistics():
    K = 1
    T = 100
    dt = 1.0
    network_hypers = {'c': np.zeros(K, dtype=np.int), 'p': 1.0, 'kappa': 10.0, 'v': 10*5.0}

    true_model = DiscreteTimeNetworkHawkesModelSpikeAndSlab(K=K, dt=dt,