import numpy as np

from conmorph.connections import RecurrentInhibition
from interneuron_polarity.model.morphology.positions import Uniform, UniformOrientation
from interneuron_polarity.model.morphology.shapes import Ellipsoid
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 somata distribution for ex and in and the axon/dendrite shape of the ins

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

polar_long_axis = 200
polar_short_axis = 40
circular_radius = (polar_long_axis*polar_short_axis**2)**(1.0/3)
volume = 4.0/3*np.pi*circular_radius**3



circular_inhibitory_shape = Ellipsoid(circular_radius, circular_radius)
polar_inhibitory_shape = Ellipsoid(polar_long_axis, polar_short_axis) #approximately same volume
inhibitory_orientation = UniformOrientation(0, np.pi / 2.0, 0, 2 * np.pi)



### Determine the excitatory neurons in one inhibitory cloud
density_ex = 500.0 / (300 ** 3)
density_in = 250.0 / (300 ** 3)
number_of_excitatory_neurons = int(density_ex * v_x * v_y * v_z)
number_of_inhibitory_neurons = int(density_in * v_x * v_y * v_z)


excitatory_somata = somata_distribution.generate(number_of_excitatory_neurons)
inhibitory_somata = somata_distribution.generate(number_of_inhibitory_neurons)
inhibitory_orientations = inhibitory_orientation.generate(number_of_inhibitory_neurons)

connectivity_p = np.zeros((number_of_inhibitory_neurons, number_of_excitatory_neurons))
connectivity_c = np.zeros((number_of_inhibitory_neurons, number_of_excitatory_neurons))


for inhibitory_index, inhibitory_specification in enumerate(zip(inhibitory_somata.T, inhibitory_orientations.T)):
    in_pos, in_orientation = inhibitory_specification
    contained_indices_p = polar_inhibitory_shape.contains(excitatory_somata, in_pos, in_orientation)
    contained_indices_c = circular_inhibitory_shape.contains(excitatory_somata, in_pos, in_orientation)
    connectivity_p[inhibitory_index, contained_indices_p] = 1
    connectivity_c[inhibitory_index, contained_indices_c] = 1

connectivities = [connectivity_p, connectivity_c]
labels=["polar", "circular"]

for conn, label in zip(connectivities, labels):
    ex_ex_conn = RecurrentInhibition(conn, excitatory_somata)
    ex_ex_conn.save("../../data/{:s}_recurrent_inhibition.pickle".format(label))