import matplotlib.pyplot as plt
import numpy as np
from brian2.units import *
from interneuron_polarity.io.io import IO


def plot_spiking(n_ex, n_in, ex_spike_trains, in_spike_trains, t_start=None, t_end=None):
    fig = plt.figure(figsize=(10, 5))
    ax = fig.add_subplot(111)
    for key, times in ex_spike_trains.items():
        ax.plot(times / ms, (n_ex + n_in - key) * np.ones(times.shape), 'r|')

    for key, times in in_spike_trains.items():
        ax.plot(times / ms, (n_in - key) * np.ones(times.shape), 'b|')

    ax.grid(axis='x')
    if t_start is not None:
        ax.set_xlim(t_start, t_end)
    ax.set_ylim(0, n_ex + n_in)
    ax.set_xlabel("Time(ms)")

    return fig


def plot_network_activity(path_to_simulation_dir):
    my_io = IO(path_to_simulation_dir)
    ex_spike_trains, in_spike_trains = my_io.load_spike_trains()

    n_ex = len(ex_spike_trains.keys())
    n_in = len(in_spike_trains.keys())

    fig = plot_spiking(n_ex, n_in, ex_spike_trains, in_spike_trains)
    my_io.save_figure(fig, "network_activity_spike_raster")
