import numpy as np
import scipy.stats as stats

from scripts.spatial_maps.spatial_network_layout import get_excitatory_phases_in_inhibitory_axon, create_grid_of_excitatory_neurons, \
    Interneuron


def distance(interneuron_1, interneuron_2):
    return np.sqrt((interneuron_1.x - interneuron_2.x) ** 2 + (interneuron_1.y - interneuron_2.y) ** 2)


def contains(interneuron, XX, YY):
    d1, d2 = distance_to_points(interneuron, XX, YY)
    return d1 + d2 < interneuron.c


def distance_to_points(interneuron, XX, YY):
    x_p_1, y_p_1 = interneuron.get_p1()
    x_p_2, y_p_2 = interneuron.get_p2()

    d1 = np.sqrt((XX - x_p_1) ** 2 + (YY - y_p_1) ** 2)
    d2 = np.sqrt((XX - x_p_2) ** 2 + (YY - y_p_2) ** 2)
    return d1, d2


def get_overlap(interneuron_1, interneuron_2, ds=0.1):
    if distance(interneuron_1, interneuron_2) > 2 * (max(interneuron_1.a, interneuron_1.b) + max(interneuron_2.a, interneuron_2.b)):
        overlap = 0
    else:
        bb_x_min = interneuron_1.x - max(interneuron_1.a, interneuron_1.b)
        bb_x_max = interneuron_1.x + max(interneuron_1.a, interneuron_1.b)
        bb_y_min = interneuron_1.y - max(interneuron_1.a, interneuron_1.b)
        bb_y_max = interneuron_1.y + max(interneuron_1.a, interneuron_1.b)
        bb_X, bb_Y = np.meshgrid(np.arange(bb_x_min, bb_x_max, ds), np.arange(bb_y_min, bb_y_max, ds))
        in_1 = contains(interneuron_1, bb_X, bb_Y)
        in_2 = contains(interneuron_2, bb_X, bb_Y)
        overlap = np.sum(np.logical_and(in_1, in_2)) * ds ** 2
    return overlap


def get_uniform_grid(sheet_size, grid_points_per_row):
    XX, YY = np.meshgrid(np.linspace(0, sheet_size, grid_points_per_row),
                         np.linspace(0, sheet_size, grid_points_per_row))
    return np.vstack((XX.reshape((np.prod(XX.shape),)), YY.reshape((np.prod(YY.shape),)))).T


def get_overlap_matrix(axonal_clouds, ds):
    overlaps = np.zeros((len(axonal_clouds), len(axonal_clouds)))

    for row_idx, cloud1 in enumerate(axonal_clouds):
        for column_idx, cloud2 in enumerate(axonal_clouds[row_idx + 1:]):
            overlaps[row_idx, row_idx + 1 + column_idx] = get_overlap(cloud1, cloud2, ds)
    return overlaps


def get_entropy(cloud, ex_positions, ex_tunings, entropy_bins):
    phase_vals = get_excitatory_phases_in_inhibitory_axon(ex_positions, ex_tunings, cloud)
    phase_distr, _ = np.histogram(phase_vals, entropy_bins, density=True)
    return stats.entropy(phase_distr)


def get_entropies(axonal_clouds, ex_positions, ex_tunings, entropy_bins):
    return [get_entropy(cloud, ex_positions, ex_tunings, entropy_bins) for cloud in axonal_clouds]


class SpatialLayout(object):
    def __init__(self, orientation_map, NE, NI, long_axis, short_axis, sheet_size):
        self.orientation_map = orientation_map
        self.NE = NE
        self.NI = NI
        self.long_axis = long_axis
        self.short_axis = short_axis
        self.sheet_size = sheet_size

        ex_positions, ex_tunings = create_grid_of_excitatory_neurons(sheet_size, int(np.sqrt(NE)),
                                                                     orientation_map)
        self.ex_positions = ex_positions
        self.ex_tunings = ex_tunings
        self.in_positions = get_uniform_grid(sheet_size, int(np.sqrt(NI)))

    def get_axonal_clouds(self, orientation_array):
        return [Interneuron(in_x_y[0], in_x_y[1], self.long_axis, self.short_axis, orientation) for in_x_y, orientation in
                zip(
                    self.in_positions, orientation_array)]

    def get_entropies(self, orientation_array, number_of_bins):
        entropy_bins = np.linspace(-np.pi, np.pi, number_of_bins + 1)
        axonal_clouds = self.get_axonal_clouds(orientation_array)
        return get_entropies(axonal_clouds, self.ex_positions, self.ex_tunings, entropy_bins)

    def get_mean_entropy(self, orientation_array, number_of_bins):
        return np.mean(self.get_entropies(orientation_array, number_of_bins))

    def get_negative_mean_entropy(self, orientation_array, number_of_bins):
        return -self.get_mean_entropy(orientation_array, number_of_bins)

    def get_overlap(self, orientation_array, ds):
        axonal_clouds = self.get_axonal_clouds(orientation_array)
        overlaps = get_overlap_matrix(axonal_clouds, ds)
        return np.sum(overlaps)