import numpy as np
import scipy.stats as stats

from scripts.interneuron_placement import get_excitatory_phases_in_inhibitory_axon, create_grid_of_excitatory_neurons, \
    Pickle


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


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


def distance_to_points(pickle, XX, YY):
    x_p_1, y_p_1 = pickle.get_p1()
    x_p_2, y_p_2 = pickle.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(pickle_1, pickle_2, ds=0.1):
    if distance(pickle_1, pickle_2) > 2 * (max(pickle_1.a, pickle_1.b) + max(pickle_2.a, pickle_2.b)):
        overlap = 0
    else:
        bb_x_min = pickle_1.x - max(pickle_1.a, pickle_1.b)
        bb_x_max = pickle_1.x + max(pickle_1.a, pickle_1.b)
        bb_y_min = pickle_1.y - max(pickle_1.a, pickle_1.b)
        bb_y_max = pickle_1.y + max(pickle_1.a, pickle_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(pickle_1, bb_X, bb_Y)
        in_2 = contains(pickle_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, 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 [Pickle(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)