import numpy
from interneuron_polarity.io.io import IO
from interneuron_polarity.model.morphology.shapes import Ellipsoid
from tqdm import tqdm


def create_connectivity(simulation_directory):
    my_io = IO(simulation_directory)
    morphology = my_io.load_morphology_dict()
    ex_positions, in_positions, in_orientations = my_io.load_positions_and_orientations()

    n_ex = ex_positions.shape[1]
    n_in = in_positions.shape[1]

    long_axis = morphology["inhibitory_neurons"]["long_axis"]
    short_axis = morphology["inhibitory_neurons"]["short_axis"]

    inhibitory_axon_shape = Ellipsoid(long_axis, short_axis)

    in_ex_connectivity = get_in_ex_connectivity(n_ex, n_in, ex_positions, in_positions, inhibitory_axon_shape,
                                                in_orientations)

    my_io.save_in_ex_connectivity(in_ex_connectivity)


def get_in_ex_connectivity(n_ex, n_in, ex_positions, in_positions, inhibitory_axon_shape, in_orientations):
    in_ex_connectivity = numpy.zeros((n_in, n_ex))
    t = tqdm(total=n_in, desc="Identify morphological connection of interneurons to excitatory neurons")

    for inhibitory_index, inhibitory_specification in enumerate(zip(in_positions.T, in_orientations.T)):
        in_pos, in_orientation = inhibitory_specification
        contained_excitatory_neurons = inhibitory_axon_shape.contains(ex_positions, in_pos,
                                                                      in_orientation)

        in_ex_connectivity[inhibitory_index, contained_excitatory_neurons] = 1
        t.update(1)
    t.close()
    return in_ex_connectivity
