import numpy as np

from interneuron_polarity.model.morphology.volume import NeuralVolume


def uniform(distance_to_foci):
    return np.ones(distance_to_foci.shape)


def get_gaussian(length_scale):
    def gaussian(distance_to_foci):
        return np.exp(-distance_to_foci ** 2 / length_scale)

    return gaussian


class Ellipsoid:
    def __init__(self, r_long, r_small):
        self.center = np.array([0, 0, 0])
        self.r_long = r_long
        self.r_small = r_small

        self.distance_between_foci = 2 * np.sqrt(r_long ** 2 - r_small ** 2)
        self.distance_to_foci = 2 * r_long

    def get_volume(self):
        return np.pi*4/3.0*self.r_small**2*self.r_long

    def get_containing_volume(self, center, long_axis, voxel_size):
        first_direction = tuple(
            center + long_axis * self.r_long)  # is this big enough?, not quite add orthogonal direction with r_small
        second_direction = tuple(center - long_axis * self.r_long)
        x_min, y_min, z_min = tuple([min(f, s) for f, s in zip(first_direction, second_direction)])
        x_max, y_max, z_max = tuple([max(f, s) for f, s in zip(first_direction, second_direction)])
        return NeuralVolume(x_max, y_max, z_max, voxel_size, x_min, y_min, z_min)

    def get_density(self, center, long_axis, volume, density=uniform):
        coordinates = volume.get_coordinates()
        distance_to_foci = self.get_distance_to_foci(coordinates, center, long_axis)

        density = np.where(distance_to_foci < self.distance_to_foci, density(distance_to_foci),
                           np.zeros(distance_to_foci.shape))
        return density.reshape(volume.mesh_shape) / np.linalg.norm(density)

    def get_distance_to_foci(self, coordinates, center, long_axis):
        focus_1 = center + long_axis * self.distance_between_foci / 2.0
        focus_2 = center - long_axis * self.distance_between_foci / 2.0
        number_of_coordinates = coordinates.shape[1]
        distance_to_focus_1 = np.linalg.norm(
            coordinates - np.repeat(np.expand_dims(focus_1, axis=1), number_of_coordinates, axis=1), axis=0)
        distance_to_focus_2 = np.linalg.norm(
            coordinates - np.repeat(np.expand_dims(focus_2, axis=1), number_of_coordinates, axis=1), axis=0)
        distance_to_foci = distance_to_focus_1 + distance_to_focus_2
        return distance_to_foci

    def contains(self, coordinates, center, long_axis):
        return np.where(self.get_distance_to_foci(coordinates, center, long_axis) < self.distance_to_foci)[0]


def get_overlap(density_1, density_2):
    return np.sum(density_1 * density_2)


def distance(x, y):
    return np.linalg.norm(x - y)


class NeuralProcess:
    def __init__(self, form, orientation):
        self.form = form
        self.orientation = orientation

