import logging

import numpy as np
from interneuron_polarity.io.io import IO, get_slice_dimensions
from interneuron_polarity.model.morphology.positions import Uniform, UniformOrientation

_logger = logging.getLogger(__name__)


def simulate_somata_distribution(density, v_x, v_y, v_z):
    somata_distribution = Uniform(np.array([0, 0, 0]), np.array([v_x, v_y, v_z]))
    number_of_neurons = int(density * v_x * v_y * v_z)
    return somata_distribution.generate(number_of_neurons)


def simulate_axonal_orientations(number_of_axons):
    orientation_distribution = UniformOrientation(0, np.pi / 2.0, 0, 2 * np.pi)
    return orientation_distribution.generate(number_of_axons)


def create_spatial_arrangement(path_to_simulation_dir):
    # instead of saving everything, one could also generate it every time by setting a random seed
    my_io = IO(path_to_simulation_dir)

    morphology_dict = my_io.load_morphology_dict()

    v_x, v_y, v_z = get_slice_dimensions(morphology_dict)

    ex_density = morphology_dict["excitatory_neurons"]["density"]
    excitatory_somata_positions = simulate_somata_distribution(ex_density, v_x, v_y, v_z)
    n_ex = excitatory_somata_positions.shape[1]
    in_density = morphology_dict["inhibitory_neurons"]["density"]
    inhibitory_somata_positions = simulate_somata_distribution(in_density, v_x, v_y, v_z)
    n_in = inhibitory_somata_positions.shape[1]

    inhibitory_axonal_orientations = simulate_axonal_orientations(n_in)

    print("Generated positions and orientations for {:d} excitatory and {:d} inhibitory neurons".format(n_ex, n_in))

    my_io.save_positions_and_orientations(excitatory_somata_positions,
                                          inhibitory_somata_positions,
                                          inhibitory_axonal_orientations)


