import numpy as np
import matplotlib.pyplot as plt

morphologies =["polar", "circular"]

image_folder="../reports/images/"
data_folder ="../../data/"


connectivity_dict={}

for c in connectivity:
    connectivity_dict[c]={}
    for morphology in morphologies:
        connectivity_dict[c][morphology] = np.load("{:s}{:s}_connectivity_{:s}.npy".format(data_folder, c, morphology))
        connectivity_dict[c]["distances"] = np.load("{:s}{:s}_connectivity_{:s}.npy".format(data_folder, c,
                                                                                        "distance_bins"))


# Plot connectivity matrix with neuron index sorted by their (x,y) coords?


fig=plt.figure()
ax = fig.add_subplot(111)

colors = ["blue", "orange"]
linestyle = ["-", "--" ]

for c, morphology in zip(connectivity, morphologies):

    distance_bins = connectivity_dict[c]["distances"]
    for morphology, color in zip(morphologies, colors):
        number_of_connections = connectivity_dict[c][morphology]
        ax.plot(distance_bins[:-1], number_of_connections, label="{:s} {:s}".format(c, morphology), color=color, ls=ls)

ax.grid()
ax.legend()

ax.set_xlabel("Distance (um)")
ax.set_ylabel("Average number of connections")
fig.savefig(image_folder+"distance_dependent_connectivity.png", dpi=300)

# Effective interaction
interaction_first_order = -1.0
interaction_second_order = -interaction_first_order/10.0


fig=plt.figure()
ax = fig.add_subplot(111)

interaction_weights = [interaction_first_order, interaction_second_order]

for morphology, color in zip(morphologies, colors):
    distance_bins = connectivity_dict[connectivity[1]]["distances"]
    effective_interaction = np.zeros((distance_bins[:-1].shape))
    first_order_interaction = interaction_weights[0]*connectivity_dict[connectivity[0]][morphology]
    second_order_interaction = interaction_weights[1]*connectivity_dict[connectivity[1]][morphology]
    effective_interaction[:first_order_interaction.shape[0]] = first_order_interaction
    effective_interaction = effective_interaction+second_order_interaction

    ax.plot(distance_bins[1:], effective_interaction, label="weight_fo {:.1f} weight_so {:.1f} {:s}".format(
        interaction_first_order, interaction_second_order, morphology),
            color=color, ls=ls)

ax.grid()
ax.legend()

ax.set_xlabel("Distance (um)")
ax.set_ylabel("Effective interaction (a.u.)")
fig.savefig(image_folder+"effective_interaction.png", dpi=300)