import neo
import numpy as np
import quantities as pq
import matplotlib.pyplot as plt

import viziphant
from elephant import asset

np.random.seed(10)
spiketrain = np.linspace(0, 50, num=10)
np.random.shuffle(spiketrain)
spiketrains = np.c_[spiketrain, spiketrain + 100]
spiketrains += np.random.random_sample(spiketrains.shape) * 5
spiketrains = [neo.SpikeTrain(st, units='ms', t_stop=1 * pq.s)
               for st in spiketrains]
asset_obj = asset.ASSET(spiketrains, bin_size=3 * pq.ms)

imat = asset_obj.intersection_matrix()
pmat = asset_obj.probability_matrix_analytical(imat,
                                               kernel_width=50 * pq.ms)
jmat = asset_obj.joint_probability_matrix(pmat, filter_shape=(5, 1),
                                          n_largest=3)
mmat = asset_obj.mask_matrices([pmat, jmat], thresholds=.9)
cmat = asset_obj.cluster_matrix_entries(mmat, max_distance=11,
                                        min_neighbors=3, stretch=5)
sses = asset_obj.extract_synchronous_events(cmat)

viziphant.asset.plot_synchronous_events(spiketrains, sse=sses[1], s=10)
plt.show()