import numpy as np
from tqdm import tqdm

from conmorph.connections import OverlapMatrix
from interneuron_polarity.model.morphology.shapes import get_overlap, distance


class NeuronType:
    def __init__(self, somata, axon, dendrite):
        self.somata = somata
        self.axon = axon
        self.dendrite = dendrite

    def generate_instance(self, number):
        somata_positions = self.somata.generate(number)
        axon_orientations = self.axon.orientation.generate(number)
        dendrite_orientations = self.dendrite.orientation.generate(number)
        return NeuronTypeInstance(self, number, somata_positions, axon_orientations, dendrite_orientations)


class NeuronTypeInstance:
    def __init__(self, type, number, somata_positions, axon_orientations, dendrite_orientations):
        self.type = type
        self.number = number
        self.somata_positions = somata_positions
        self.axon_orientations = axon_orientations
        self.dendrite_orientations = dendrite_orientations

    def overlap_onto(self, other_instance, V):
        overlap = np.zeros((self.number, other_instance.number))

        # calculate all overlaps of this axons with other dendrites
        number_of_overlaps = self.number*other_instance.number
        t = tqdm(total=number_of_overlaps)
        #!!! Make this faster by just considering the volume around the axon
        # Waste less volume by getting the minimal rectangle around the form and then turn the orientation of the other form
        for projecting_idx in np.arange(0, self.number):
            center_projecting = self.somata_positions[:, projecting_idx]
            orientation = self.axon_orientations[:, projecting_idx]
            volume = V
            density_axon = self.type.axon.form.get_density(center_projecting, orientation, volume)
            for projected_idx in np.arange(0, other_instance.number):
                center_projected = other_instance.somata_positions[:, projected_idx]

                if distance(center_projecting, center_projected) > (self.type.axon.form.r_long+other_instance.type.dendrite.form.r_long):
                    overlap[projecting_idx, projected_idx] = 0
                else:
                    density_dendrite = other_instance.type.dendrite.form.get_density(center_projected, other_instance.dendrite_orientations[:, projected_idx], volume)

                    overlap[projecting_idx, projected_idx] = get_overlap(density_axon,density_dendrite)
                t.update(1)

        return overlap


class NeuralTissue(object):
    def __init__(self, volume, excitatory_population, inhibitory_population):
        self.volume = volume
        self.excitatory_population = excitatory_population
        self.inhibitory_population = inhibitory_population

    def calculate_overlap_matrix(self, number_of_excitatory_neurons, number_of_inhibitory_neurons):
        total_number_of_neurons = number_of_excitatory_neurons+number_of_inhibitory_neurons
        excitatory_neurons = self.excitatory_population.generate_instance(number_of_excitatory_neurons)
        inhibitory_neurons = self.inhibitory_population.generate_instance(number_of_inhibitory_neurons)

        overlaps = np.zeros((total_number_of_neurons, total_number_of_neurons))

        print("Ex-ex")
        overlaps[:number_of_excitatory_neurons, :number_of_excitatory_neurons] = excitatory_neurons.overlap_onto(excitatory_neurons, self.volume)
        print("Ex-in")
        overlaps[:number_of_excitatory_neurons, number_of_excitatory_neurons:] = excitatory_neurons.overlap_onto(inhibitory_neurons, self.volume)
        print("In-ex")
        overlaps[number_of_excitatory_neurons:, :number_of_excitatory_neurons] = inhibitory_neurons.overlap_onto(excitatory_neurons, self.volume)
        print("In-in")
        overlaps[number_of_excitatory_neurons:, number_of_excitatory_neurons:] = inhibitory_neurons.overlap_onto(inhibitory_neurons, self.volume)

        overlap = OverlapMatrix(excitatory_neurons, inhibitory_neurons, overlaps)

        return overlap