import scipy as sp
import scipy.optimize
import scipy.spatial

from .network import EmbeddedNetwork
from modelflow import Node, Input
from .bases import NetworkModelMixin
from .util import diag_to_zero


class EXP(NetworkModelMixin, Node):
    """
    Network with exponentially decaying connectivity (liquid state machine).

    :param decay: Float in [0, 1]. If set to 1, the decay is strongest. If set to 0 there is no decay at all.
    """
    GUESSED_ALPHA_EXC = 0.49
    GUESSED_ALPHA_INH = 1.25
    SPATIAL_DIMENSION = 3
    decay = Input()

    @staticmethod
    def connection_probability(distance, alpha, p_0):
        return sp.exp(- distance / alpha) * p_0

    @staticmethod
    def p_0(decay, target_connectivity):
        """
        Connectivty :math:`p_0` for zero distance as function
        of the decay strengh and the target connectivity.
        The connectivity interpolates :math:`p_0` linearly between
        the target connectivity and 1, according to ``decay_strength``.

        Parameters
        ----------
        decay: float in [0,1]
            How fast the connection probability should decay with distance.
            A value of 1 indicates fastest decay, 0 indicates no decay et all.

        target_connectivity: float in [0, 1]
            Desired overall target connectivity.

        Returns
        -------
        p_0: float in [0, 1]
            Connectivity for zero distance.
        """
        return target_connectivity + (1 - target_connectivity) * decay

    @property
    def p_0_exc(self):
        return self.p_0(self.decay, self.p_exc)

    @property
    def p_0_inh(self):
        return self.p_0(self.decay, self.p_inh)

    def run(self) -> EmbeddedNetwork:
        positions = self._generate_positions()
        pairwise_distances = self._generate_pairwise_distances(positions)
        alpha_exc, alpha_inh = self._generate_alphas(pairwise_distances)
        self.alpha_exc_, self.alpha_inh_ = alpha_exc, alpha_inh
        adjacency_matrix = self._adjacency_matrix_from_distances(alpha_exc, alpha_inh, pairwise_distances)
        net = self._adjacency_matrix_to_network(adjacency_matrix, positions)
        return net

    def _adjacency_matrix_to_network(self, W, positions):
        net = EmbeddedNetwork(W, nr_exc=self.nr_exc, nr_inh=self.nr_inh)
        net.positions = positions
        net.random_relabel_neurons()
        return net

    def _adjacency_matrix_from_distances(self, alpha_exc, alpha_inh, pairwise_distances):
        probabilities = sp.empty(pairwise_distances.shape)
        probabilities[:, :self.nr_exc] = self.connection_probability(pairwise_distances[:, :self.nr_exc], alpha_exc, self.p_0_exc)
        probabilities[:, self.nr_exc:] = self.connection_probability(pairwise_distances[:, self.nr_exc:], alpha_inh, self.p_0_inh)
        W = (probabilities > sp.rand(*probabilities.shape)).astype(float)
        W[:, self.nr_exc:] *= -1
        diag_to_zero(W)
        return W

    def _generate_alphas(self, pairwise_distances):
        alpha_exc = self._p_to_alpha(self.p_exc, pairwise_distances, self.p_0_exc)
        alpha_inh = self._p_to_alpha(self.p_inh, pairwise_distances, self.p_0_inh)
        return alpha_exc, alpha_inh

    def _generate_positions(self):
        return sp.rand(self.nr_neurons, self.SPATIAL_DIMENSION)

    @staticmethod
    def _generate_pairwise_distances(positions):
        return sp.spatial.distance.squareform(sp.spatial.distance.pdist(positions))

    def _alpha_to_p(self, alpha: float, pairwise_distances, p_0) -> float:
        neuron_to_neuron_connection_probability = self.connection_probability(pairwise_distances, alpha, p_0)
        diag_to_zero(neuron_to_neuron_connection_probability)
        average_connectivity = neuron_to_neuron_connection_probability.sum() / (self.nr_neurons * (self.nr_neurons - 1))
        return average_connectivity

    def _p_to_alpha(self, p: float, pairwise_distances, p_0) -> float:
        p_threshold = (.3 + .6) / 2  # assume the default connectivityes of 0.3 and 0.6
        initial_alpha = self.GUESSED_ALPHA_EXC if p < p_threshold else self.GUESSED_ALPHA_INH
        return sp.optimize.newton(lambda alpha: self._alpha_to_p(alpha, pairwise_distances, p_0) - p, initial_alpha, tol=1e-4)
