import matplotlib.pyplot as plt
import pandas
from pypet import Trajectory
from pypet.brian2 import Brian2MonitorResult

from scripts.competition_gain.run_micro_network import TRAJ_NAME, UNI_DIRECTIONAL, RECIPROCAL, NO_CONNECTIONS
from scripts.spatial_network.run_entropy_maximisation_orientation_map import DATA_FOLDER

if __name__ == "__main__":
    traj = Trajectory(TRAJ_NAME, add_time=False, dynamic_imports=Brian2MonitorResult)
    traj.f_load(filename=DATA_FOLDER + TRAJ_NAME + ".hdf5", load_parameters=2, load_results=2)

    firing_rates_frame = pandas.DataFrame(traj.results.summary)
    firing_rates_frame["delta_f"] = firing_rates_frame["f_1"] - firing_rates_frame["f_2"]

    fig, ax = plt.subplots(1, 1, figsize=(3, 4))
    colors = {NO_CONNECTIONS: "grey", RECIPROCAL: "blue", UNI_DIRECTIONAL: "orange"}

    labels = {
        NO_CONNECTIONS: "no inh",
        RECIPROCAL: "hub inh",
        UNI_DIRECTIONAL: "relay inh"
    }

    for architecture in [UNI_DIRECTIONAL, RECIPROCAL, NO_CONNECTIONS]:
        df = firing_rates_frame[firing_rates_frame["architecture"] == architecture]
        ax.plot(df["I_1"].values, df["delta_f"].values, label=labels[architecture], color=colors[architecture], lw=1.5,
                zorder=10)
        ax.plot(df["I_1"].values, df["f_1"].values, linestyle='--', color=colors[architecture], alpha=0.6, zorder=5)
        ax.plot(df["I_1"].values, df["f_2"].values, linestyle=':', color=colors[architecture], alpha=0.6, zorder=5)
    ax.set_xlabel("I$_1$ - I$_2$ (nA)")
    ax.set_ylabel("firing rate (Hz)")
    ax.legend()
    ax.text(0.1, 55, "f$_1$", ha="right", va="bottom")
    ax.text(0.23, 10, "f$_2$", ha="right", va="top")
    ax.text(0.55, 75, "f$_1$-f$_2$", ha="right", va="top")

    fig.tight_layout()
    fig.savefig("../../figures/competition_gain/rate_gain.png", dpi=300)

    N_E = traj.par.network.N_E
    N_I = traj.par.network.N_I

    ex_index_mappings = {
        RECIPROCAL: (lambda idx: N_E + N_I + 2 - idx if idx < int(N_E / 2.0) else N_E - idx),
        UNI_DIRECTIONAL: (lambda idx: N_E + 2 * N_I + 2 - idx if idx < int(N_E / 2.0) else N_E - idx)
    }

    in_index_mappings = {
        RECIPROCAL: (lambda idx: int(N_E / 2.0) + 2 + idx),
        UNI_DIRECTIONAL: (lambda idx: N_E + 2 * N_I + 2 - int(N_E / 2.0) - idx)
    }

    for input_to_1 in firing_rates_frame["I_1"].values:
        df = firing_rates_frame[
            (firing_rates_frame["I_1"] == input_to_1) & (firing_rates_frame["architecture"] != NO_CONNECTIONS)]

        fig, axes = plt.subplots(2, 1, sharex=True)
        for architecture, run_idx in zip(df["architecture"].values, df["run"].values):
            ax_idx = 0 if architecture == RECIPROCAL else 1
            ax = axes[ax_idx]

            traj.v_idx = run_idx

            neuron_indices = traj.results.crun.spikes.e.i
            spike_times = traj.results.crun.spikes.e.t
            ax.plot(spike_times*1000, list(map(ex_index_mappings[architecture], neuron_indices)), 'r|')

            if traj.results.crun.spikes.i.num_spikes > 0:
                neuron_indices = traj.results.crun.spikes.i.i
                spike_times = traj.results.crun.spikes.i.t

                ax.plot(spike_times*1000, list(map(in_index_mappings[architecture], neuron_indices)), 'b|')
            ax.set_xlim(0,100)
            ax.yaxis.set_ticks([])
        axes[1].set_xlabel("t (ms)")
        fig.suptitle("I$_1$-I$_2$ = {:.1f} nA".format(input_to_1))

        fig.savefig("../../figures/competition_gain/spike_raster_delta_I_{:.1f}.png".format(input_to_1), dpi=300)