import numpy as np

from conmorph.connections import save_overlap
from conmorph.neurons import NeuronType, NeuralTissue
from interneuron_polarity.model.morphology.positions import Uniform, Deterministic, UniformOrientation
from interneuron_polarity.model.morphology.shapes import Ellipsoid, NeuralProcess
from interneuron_polarity.model.morphology.volume import NeuralVolume

### Define the neural volume
voxel_size = 10
v_x = 300
v_y = 300
v_z = 300

V = NeuralVolume(v_x, v_y, v_z, voxel_size)

### Define neural types with their somata distribution and the axon and dendritic tree shapes

somata_distribution = Uniform(np.array([0,0,0]), np.array([v_x, v_y, v_z]))

excitatory_axon_shape = Ellipsoid(100, 100)
excitatory_axon_orientation = Deterministic(np.array([0, 0, 1]))
excitatory_axon = NeuralProcess(excitatory_axon_shape, excitatory_axon_orientation)

excitatory_dendrite_shape = Ellipsoid(100, 100)
excitatory_dendrite_orientation = Deterministic(np.array([0, 0, 1]))
excitatory_dendrite = NeuralProcess(excitatory_dendrite_shape, excitatory_dendrite_orientation)

circular_inhibitory_axon_shape = Ellipsoid(80, 80)
polar_inhibitory_axon_shape = Ellipsoid(200, 50) #approximately same volume
inhibitory_axon_orientation = UniformOrientation(0, np.pi/2.0, 0, 2*np.pi)
circular_inhibitory_axon = NeuralProcess(circular_inhibitory_axon_shape, inhibitory_axon_orientation)
polar_inhibitory_axon = NeuralProcess(polar_inhibitory_axon_shape, inhibitory_axon_orientation)

inhibitory_dendrite_shape = Ellipsoid(100, 100)
inhibitory_dendrite_orientation = Deterministic(np.array([0, 0, 1]))
inhibitory_dendrite = NeuralProcess(inhibitory_dendrite_shape, inhibitory_dendrite_orientation)

excitatory_population = NeuronType(somata_distribution, excitatory_axon, excitatory_dendrite)
inhibitory_populations = [NeuronType(somata_distribution, inhibitory_axon, inhibitory_dendrite) for inhibitory_axon in [circular_inhibitory_axon, polar_inhibitory_axon]]


### Calculate the axon-dendrite overlap between all pairs of neurons
number_of_excitatory_neurons = 500
number_of_inhibitory_neurons = 250

tissues = [ NeuralTissue(V, excitatory_population, inhibitory_population) for inhibitory_population in inhibitory_populations]

file_path = "/home/pfeiffer/Projects/subiculum_interneuron_polarity/data"


number_of_overlap_matrices = 2

for tissue, type in zip(tissues, ["circular_inhibitory_neurons", "polar_inhibitory_neurons"]):
    for idx in range(number_of_overlap_matrices):
        print("Calculate {:s} overlap matrix {:d}/{:d}".format(type, idx+1, number_of_overlap_matrices))
        file_name = "{:s}_overlap_matrix_{:d}.pickle".format(type, idx)
        path_to_store = file_path + "/" + file_name

        overlap_matrix = tissue.calculate_overlap_matrix(number_of_excitatory_neurons, number_of_inhibitory_neurons)

        save_overlap(path_to_store, overlap_matrix)