import neo
import numpy as np
import quantities as pq
import matplotlib.pyplot as plt
from elephant.spike_train_generation import homogeneous_poisson_process
from viziphant.rasterplot import eventplot
from viziphant.events import add_event
np.random.seed(12)
spiketrains = [homogeneous_poisson_process(rate=5*pq.Hz,
               t_stop=10*pq.s) for _ in range(20)]

fig, axes = plt.subplots(2, 2, sharex=True, sharey='row')
event = neo.Event([0.5, 8]*pq.s, labels=['trig0', 'trig1'])
eventplot(spiketrains[:10], axes=axes[:, 0], histogram_bins=20,
          title="Neuron A", linelengths=0.75, linewidths=1)
add_event(axes[:, 0], event)
eventplot(spiketrains[10:], axes=axes[:, 1], histogram_bins=20,
          title="Neuron B", linelengths=0.75, linewidths=1)
add_event(axes[:, 1], event)
plt.show()