# Path: pyhawkes/internals/network.py
# class StochasticBlockModel(GibbsSBM, MeanFieldSBM):
#     pass
# 
# 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 adjusted_mutual_info_score, adjusted_rand_score
from pyhawkes.internals.network import StochasticBlockModel
from pyhawkes.models import DiscreteTimeNetworkHawkesModelSpikeAndSlab
from pybasicbayes.util.text import progprint_xrange
        seed = np.random.randint(2**32)

    print("Setting seed to ", seed)
    np.random.seed(seed)

    C = 2
    K = 100
    c = np.arange(C).repeat(np.ceil(K/float(C)))[:K]
    T = 1000
    dt = 1.0
    B = 3

    # Generate from a true model
    true_p = np.random.rand(C,C) * 0.25
    true_network = StochasticBlockModel(K, C, c=c, p=true_p, v=10.0)
    true_model = \
        DiscreteTimeNetworkHawkesModelSpikeAndSlab(
                K=K, dt=dt, B=B, network=true_network)

    S,R = true_model.generate(T)

    # Plot the true network
    plt.ion()
    true_im = true_model.plot_adjacency_matrix()
    plt.pause(0.001)


    # Make a new model for inference
    test_network = StochasticBlockModel(K, C, beta=1./K)
    test_model = \

# Path: pyhawkes/internals/network.py
# class StochasticBlockModel(GibbsSBM, MeanFieldSBM):
#     pass
# 
# 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_sbm_gibbs.py
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import adjusted_mutual_info_score, adjusted_rand_score
from pyhawkes.internals.network import StochasticBlockModel
from pyhawkes.models import DiscreteTimeNetworkHawkesModelSpikeAndSlab
from pybasicbayes.util.text import progprint_xrange

        seed = np.random.randint(2**32)

    print("Setting seed to ", seed)
    np.random.seed(seed)

    C = 2
    K = 100
    c = np.arange(C).repeat(np.ceil(K/float(C)))[:K]
    T = 1000
    dt = 1.0
    B = 3

    # Generate from a true model
    true_p = np.random.rand(C,C) * 0.25
    true_network = StochasticBlockModel(K, C, c=c, p=true_p, v=10.0)
    true_model = \
        DiscreteTimeNetworkHawkesModelSpikeAndSlab(
                K=K, dt=dt, B=B, network=true_network)

    S,R = true_model.generate(T)

    # Plot the true network
    plt.ion()
    true_im = true_model.plot_adjacency_matrix()
    plt.pause(0.001)


    # Make a new model for inference
    test_network = StochasticBlockModel(K, C, beta=1./K)
    test_model = \

        DiscreteTimeNetworkHawkesModelSpikeAndSlab(