import scipy as sp
from brian2 import ms
from matplotlib import pyplot as plt


def pool_scatter(pool_activations, t_list, visualization_power=1):
    fig, ax = plt.subplots()
    max_active_nr = pool_activations.max()
    for k, pool in enumerate(pool_activations):
        spike_indices = sp.nonzero(pool)
        spike_t = t_list[spike_indices] / ms
        y = sp.ones(len(spike_t)) * k
        sizes = pool[spike_indices]**visualization_power
        sizes /= max_active_nr**visualization_power
        sizes *= 40
        ax.scatter(spike_t, y, s=sizes, c="k")

    ax.set_xlabel("t (ms)")
    ax.set_ylabel("Pool nr")
    ax.set_title("Pool activations. Size relative to max activation, "
                 "power={}".format(visualization_power), y = 1.05)
    return fig, ax