# 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
import matplotlib.pyplot as plt
from sklearn.metrics import roc_auc_score
from pyhawkes.models import DiscreteTimeNetworkHawkesModelSpikeAndSlab

np.random.seed(0)

def demo(K=3, T=1000, dt_max=20, p=0.25):
    """

    :param K:       Number of nodes
    :param T:       Number of time bins to simulate
    :param dt_max:  Number of future time bins an event can influence
    :param p:       Sparsity of network
    :return:
    """
    ###########################################################
    # Generate synthetic data
    ###########################################################
    network_hypers = {"p": p, "allow_self_connections": False}
    true_model = DiscreteTimeNetworkHawkesModelSpikeAndSlab(
        K=K, dt_max=dt_max,
        network_hypers=network_hypers)
    assert true_model.check_stability()

    # Sample from the true model
    S,R = true_model.generate(T=T, keep=True, print_interval=50)

    plt.ion()
    true_figure, _ = true_model.plot(color="#377eb8", T_slice=(0,100))

    ###########################################################
    # Create a test spike and slab model
    ###########################################################

# 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: examples/discrete_demo.py
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import roc_auc_score
from pyhawkes.models import DiscreteTimeNetworkHawkesModelSpikeAndSlab


np.random.seed(0)

def demo(K=3, T=1000, dt_max=20, p=0.25):
    """

    :param K:       Number of nodes
    :param T:       Number of time bins to simulate
    :param dt_max:  Number of future time bins an event can influence
    :param p:       Sparsity of network
    :return:
    """
    ###########################################################
    # Generate synthetic data
    ###########################################################
    network_hypers = {"p": p, "allow_self_connections": False}
    true_model = DiscreteTimeNetworkHawkesModelSpikeAndSlab(
        K=K, dt_max=dt_max,
        network_hypers=network_hypers)
    assert true_model.check_stability()

    # Sample from the true model
    S,R = true_model.generate(T=T, keep=True, print_interval=50)

    plt.ion()
    true_figure, _ = true_model.plot(color="#377eb8", T_slice=(0,100))

    ###########################################################
    # Create a test spike and slab model
    ###########################################################

    test_model = DiscreteTimeNetworkHawkesModelSpikeAndSlab(