import matplotlib.pyplot as plt
import numpy
import scipy
from interneuron_polarity.io.io import IO, get_slice_dimensions
# noinspection PyUnresolvedReferences
from interneuron_polarity.model.morphology.shapes import Ellipsoid
from mpl_toolkits.mplot3d import Axes3D


def add_somata(ax, position, plot_config):
    if len(position.shape) == 1:
        position = numpy.expand_dims(position, 1)
    ax.scatter(position[0, :], position[1, :], position[2, :], **plot_config)


def plot_volume_with_cell_bodies(ex_positions, in_positions, v_x, v_y, v_z):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    soma_position = [ex_positions, in_positions]
    markers = ["^", "o"]
    colors = ["blue", "red"]

    for marker, color, position in zip(markers, colors, soma_position):
        ax.scatter(position[0, :], position[1, :], position[2, :], marker=marker, color=color)

    ax.set_xlabel("um")
    ax.set_ylabel("um")
    ax.set_zlabel("um")

    ax.set_xlim(0, v_x)
    ax.set_ylim(0, v_y)
    ax.set_zlim(0, v_z)
    return fig


def plot_volume_with_connectivity(ex_positions, in_positions, in_ex_connectivity, in_idx, v_x, v_y, v_z):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    # Plot all excitatory somata
    add_somata(ax, ex_positions, {"color": "r", "marker": "^", "alpha": 0.2})

    # Plot inhibitory neuron and all connected excitatory cells
    add_somata(ax, in_positions[:, in_idx], {"color": 'g', "marker": "o", "s": 100})
    connected_somata = ex_positions[:, numpy.where(in_ex_connectivity[in_idx])[0]]
    add_somata(ax, connected_somata, {"color": 'g', "marker": "^"})

    ax.set_xlabel("um")
    ax.set_ylabel("um")
    ax.set_zlabel("um")

    ax.set_xlim(0, v_x)
    ax.set_ylim(0, v_y)
    ax.set_zlim(0, v_z)

    return fig


def plot_connectivity_number_of_targeted_excitatory_neurons(in_ex_connectivity, morpohology):
    number_of_targeted_excitatory_neurons = numpy.sum(in_ex_connectivity, axis=1)

    density_excitatory_neurons = morpohology["excitatory_neurons"]["density"]
    r_long = morpohology["inhibitory_neurons"]["long_axis"]
    r_small = morpohology["inhibitory_neurons"]["short_axis"]
    volume_inhibitory_axon = Ellipsoid(r_long, r_small).get_volume()

    expected_mean = density_excitatory_neurons*volume_inhibitory_axon
    actual_mean = numpy.mean(number_of_targeted_excitatory_neurons)
    n_ex = in_ex_connectivity.shape[1]

    prevalence, bin_edges = scipy.histogram(number_of_targeted_excitatory_neurons, bins=100, range=[0, n_ex])

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.bar(bin_edges[1:], prevalence)

    ax.axvline(expected_mean, color='k', ls='-', label="expected mean=density*volume_axon")
    ax.axvline(actual_mean, color='k', ls='--',label="actual_mean")

    ax.set_xlabel("Number of target excitatory neurons")
    ax.set_ylabel("Abundance")
    ax.set_xlim(0, n_ex)

    ax.legend()
    return fig


def plot_spatial_configuration(path_to_simulation_dir):
    my_io = IO(path_to_simulation_dir)
    morphology = my_io.load_morphology_dict()
    ex_positions, in_positions, _ = my_io.load_positions_and_orientations()
    in_ex_connectivity = my_io.load_in_ex_connectivity()

    v_x, v_y, v_z = get_slice_dimensions(morphology)

    figure_volume_with_cell_bodies = plot_volume_with_cell_bodies(ex_positions, in_positions, v_x, v_y, v_z)
    my_io.save_figure(figure_volume_with_cell_bodies, "volume_with_cell_bodies")

    for interneuron_idx in range(4):
        figure_volume_with_cell_connections = plot_volume_with_connectivity(ex_positions, in_positions,
                                                                            in_ex_connectivity, interneuron_idx, v_x,
                                                                            v_y, v_z)
        my_io.save_figure(figure_volume_with_cell_connections,
                          "volume_with_cells_connected_to_interneuron_{:d}".format(interneuron_idx))

    figure_connectivity_number_of_target_excitatory_neurons = \
        plot_connectivity_number_of_targeted_excitatory_neurons(in_ex_connectivity, morphology)
    my_io.save_figure(figure_connectivity_number_of_target_excitatory_neurons,
                      "connectivity_number_of_targeted_excitatory_neurons")
