import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_context("talk")

morphologies =["polar", "circular"]

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

connectivity = ["first_order", "second_order"]

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 average number of connections versus distance

cmperinch = 2.54

figsize = (10.0/cmperinch, 10.0/cmperinch)
fig=plt.figure(figsize=figsize)
ax = fig.add_subplot(111)


# First order

rx, ry = 3., 1.
area = rx * ry * np.pi
theta = np.arange(0, 2 * np.pi + 0.11, 0.1)
ellipsoid_verts = np.column_stack([rx / area * np.cos(theta), ry / area * np.sin(theta)])

rx= 1
ry = rx
area = rx * ry * np.pi
circular_verts = np.column_stack([rx / area * np.cos(theta), ry / area * np.sin(theta)])



c = "first_order"
distance_bins = connectivity_dict[c]["distances"]
styles = [
    {
      "label": "polar",
        "marker": ellipsoid_verts
    },
{
      "label": "circular",
    "marker": circular_verts
    }
]



for morphology, style in zip(reversed(morphologies), reversed(styles)):
    number_of_connections = connectivity_dict[c][morphology]
    ax.plot(distance_bins[:-1], number_of_connections, ls ='--')
    ax.scatter(distance_bins[0:-2:5], number_of_connections[0:-1:5], **style)

ax.grid()
ax.set_axisbelow(True)
ax.legend()

ax.set_xlim(0,400)
ax.set_xlabel("Distance (um)")
ax.set_ylabel("<# of connections>")
fig.tight_layout()
fig.savefig(image_folder+"poster_distance_dependent_first_order_connectivity.png", dpi=300)
fig.savefig(image_folder+"poster_distance_dependent_first_order_connectivity.pdf")


figsize = (10.0/cmperinch, 10.0/cmperinch)
fig=plt.figure(figsize=figsize)
ax = fig.add_subplot(111)


# First order

rx, ry = 3., 1.
area = rx * ry * np.pi
theta = np.arange(0, 2 * np.pi + 0.11, 0.1)
ellipsoid_verts = np.column_stack([rx / area * np.cos(theta), ry / area * np.sin(theta)])

rx= 1
ry = rx
area = rx * ry * np.pi
circular_verts = np.column_stack([rx / area * np.cos(theta), ry / area * np.sin(theta)])



c = "second_order"
distance_bins = connectivity_dict[c]["distances"]
styles = [
    {
      "label": "polar",
        "marker": ellipsoid_verts
    },
{
      "label": "circular",
    "marker": circular_verts
    }
]



for morphology, style in zip(reversed(morphologies), reversed(styles)):
    number_of_connections = connectivity_dict[c][morphology]
    ax.plot(distance_bins[:-1], number_of_connections, ls ='--')
    ax.scatter(distance_bins[0:-2:5], number_of_connections[0:-1:5], **style)

ax.grid()
ax.set_axisbelow(True)
ax.legend()

ax.set_xlabel("Distance (um)")
ax.set_ylabel("<# of connections>")

ax.set_xlim(0,400)
fig.tight_layout()
fig.savefig(image_folder+"poster_distance_dependent_second_order_connectivity.png", dpi=300)
fig.savefig(image_folder+"poster_distance_dependent_second_order_connectivity.pdf")