import pickle
import graph_tool as gt
import itertools
import numpy as np

from interneuron_polarity.model.morphology.volume import get_distance_from_to


class RecurrentInhibition:
    INTERNEURON_IDENTITY = "interneuron_idx"
    POSITION = "position"
    def __init__(self, in_ex_connectivity, ex_somata_positions):
        n_in, n_ex = tuple(in_ex_connectivity.shape)

        g = gt.Graph(directed=False)
        interneuron_identity = g.new_edge_property("int")
        positions = g.new_vertex_property("vector<float>")
        exs = g.add_vertex(n_ex)
        for ex in exs:
            positions[ex] = ex_somata_positions[:, g.vertex_index[ex]]

        for interneuron_idx in range(n_in):
            connected_exs_indices = np.where(in_ex_connectivity[interneuron_idx, :] == 1)[0]
            for comb in itertools.combinations(set(connected_exs_indices), 2):
                ex1, ex2 = comb
                new_edge = g.add_edge(ex1, ex2)
                interneuron_identity[new_edge] = interneuron_idx

        g.edge_properties[RecurrentInhibition.INTERNEURON_IDENTITY] = interneuron_identity
        g.vertex_properties[RecurrentInhibition.POSITION] = positions

        self.n_in = n_in
        self.n_ex = n_ex
        self.positions = ex_somata_positions
        self.g = g

    def get_distance_to_neighbors(self, excitatory_neuron_index):
        neighbors = self.g.get_out_neighbors(excitatory_neuron_index)
        distances = self.get_distance(excitatory_neuron_index, neighbors)
        return distances[0,:]

    def get_distances_to_neighbors(self):
        distances = [self.get_distance_to_neighbors(idx) for idx in np.arange(0, self.n_ex)]
        return np.concatenate(tuple(distances))

    def get_number_of_neighbors(self):
        return np.array([len(self.g.get_out_neighbors(v)) for v in self.g.vertices()])

    def get_second_order_neighbors(self, tested_vertex):
        first_order_neighbors = self.g.get_out_neighbors(tested_vertex)
        excluded_neighbors = set(first_order_neighbors).union([tested_vertex])
        second_order_neighbors = []
        for first_order_neighbor in first_order_neighbors:
            potential_second_order_neighbors = set(self.g.get_out_neighbors(first_order_neighbor))
            second_order_neighbors.extend(potential_second_order_neighbors.difference(excluded_neighbors))
        return second_order_neighbors

    def get_distance_to_second_order_neighbors(self, excitatory_idx):
        second_order_neighbors = self.get_second_order_neighbors(excitatory_idx)
        distances = self.get_distance(excitatory_idx, second_order_neighbors)
        return distances[0,:]

    def get_distances_to_second_order_neighbors(self):
        distances = [self.get_distance_to_second_order_neighbors(idx) for idx in np.arange(0, self.n_ex)]
        return np.concatenate(tuple(distances))

    def get_connected_interneuron_indices(self, excitatory_neuron_index):
        edges = self.g.get_out_edges(excitatory_neuron_index)
        return np.array([self.g.edge_properties[RecurrentInhibition.INTERNEURON_IDENTITY][edge] for edge in edges])


    def get_numbers_of_containing_axonal_clouds(self):
        number_of_containing_axonal_clouds = [len(np.unique(self.get_connected_interneuron_indices(idx)))  for idx in range(self.n_ex)]
        return number_of_containing_axonal_clouds

    def get_distance(self, from_idx, to):
        this_position = self.positions[:, from_idx]
        neighbor_positions = self.positions[:, to]
        distances = get_distance_from_to(this_position, neighbor_positions)
        return distances

    def save(self, filename):
        with open(filename, 'wb') as file:
            pickle.dump(self, file)

    @staticmethod
    def load(filename):
        with open(filename, 'rb') as file:
            recurrent_inhibition_graph = pickle.load(file)
        return recurrent_inhibition_graph




class OverlapMatrix:
    def __init__(self, excitatory_neurons, inhibitory_neurons, overlaps):
        self.excitatory_neurons = excitatory_neurons
        self.inhibitory_neurons = inhibitory_neurons
        self.overlaps = overlaps

    def ex_on_ex(self):
        return self.overlaps[:self.excitatory_neurons.number, :self.excitatory_neurons.number]

    def ex_on_in(self):
        return self.overlaps[:self.excitatory_neurons.number, self.excitatory_neurons.number:-1]

    def in_on_in(self):
        return self.overlaps[self.excitatory_neurons.number:-1, self.excitatory_neurons.number:-1]

    def in_on_ex(self):
        return self.overlaps[self.excitatory_neurons.number:-1, :self.excitatory_neurons.number]

    def get_submatrices(self):
        return self.ex_on_ex(), self.ex_on_in(), self.in_on_in(), self.in_on_ex()

def save_overlap(filename, overlap_matrix):
    with open(filename, 'wb') as file:
        pickle.dump(overlap_matrix,file)

def load_overlap(filename):
    with open(filename, 'rb') as file:
        overlap_matrix = pickle.load(file)
    return overlap_matrix