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

import viziphant
from elephant.spike_train_generation import homogeneous_poisson_process
from elephant.unitary_event_analysis import jointJ_window_analysis

np.random.seed(10)

spiketrains1 = [homogeneous_poisson_process(rate=20 * pq.Hz,
                t_stop=2 * pq.s) for _ in range(5)]
spiketrains2 = [homogeneous_poisson_process(rate=50 * pq.Hz,
                t_stop=2 * pq.s) for _ in range(5)]

spiketrains = np.stack((spiketrains1, spiketrains2), axis=1)
ue_dict = jointJ_window_analysis(spiketrains,
                                 bin_size=5 * pq.ms,
                                 win_size=100 * pq.ms,
                                 win_step=10 * pq.ms)
viziphant.unitary_event_analysis.plot_ue(spiketrains, Js_dict=ue_dict,
                                         significance_level=0.34,
                                         unit_real_ids=['1', '2'])
plt.show()