import abc

import scipy as sp
import scipy.optimize

from .util import diag_to_zero
from modelflow import Node
from modelflow.elementarynodes import Input
from .bases import NetworkModelMixin, sample_weighted_network_from_probabilities
from .featurevector import uniform_spherical_feature_fectors
import numpy as np


def feature_vector_correlations(feature_vectors):
    correlations = feature_vectors.T @ feature_vectors
    diag_to_zero(correlations)
    return correlations


def invert_sign_of_inh_to_exc_correlations(correlations, nr_inh):
    """
    Invert sign of I -> E and I -> I correlations.
    Strictly speaking no correlations any more.
    Correlation matrix not symmetric any more.
    """
    signed_correlations = correlations.copy()
    signed_correlations[:, -nr_inh:] *= -1
    return signed_correlations


class BinomialConnectivity:
    def __init__(self, p_binomial: np.ndarray):
        """

        Parameters
        ----------
        p_binomial: 2D array
            Binomial p array

        Returns
        -------

        """
        self.p_binomial = p_binomial
        self.n_binomial_fitted_ = None

    def fit(self, p_target: float):
        """
        Find the right n to yield the desired probability.

        Parameters
        ----------

        p_target:
            is usually p_exc or p_inh
            Fits the binomial n , such that the average population probabilities match p_target

        """
        def relative_connectivity_offset(n_binomial):
            probabilities = self._binomial_connection_probabilities(n_binomial)
            return probabilities.sum() / (probabilities.shape[0]-1) / probabilities.shape[1] / p_target - 1
        self.n_binomial_fitted_ = float(sp.optimize.fsolve(relative_connectivity_offset, 1.))

    def _binomial_connection_probabilities(self, n_binomial):
        """
        Takes the correlation to a power and takes care of negative correlations.
        Takes into account whether pre/postsyn partner
        is exc or inh.

        Returns
        -------

        A matrix of connection probabilities.
        """
        return 1 - (1 - self.p_binomial)**n_binomial

    def predict(self):
        """
        Returns
        -------
        ndarray of fitted  probabilities
        """
        return self._binomial_connection_probabilities(self.n_binomial_fitted_)


class NegativeCorrelationRelaxingNetwork(NetworkModelMixin, Node):
    n_pow = Input()
    feature_space_dimension = Input()

    @abc.abstractmethod
    def signed_to_positive_correlations(self, signed_correlations: np.ndarray) -> np.ndarray:
        """"
        Define here how to relax the correlations to obain all positive correlations.
        This is a template method.
        """

    def positive_correlations_to_probability(self, positive_correlations: np.ndarray) -> np.ndarray:
        """

        Parameters
        ----------
        positive_correlations: np.ndarray
            Relaxed correlations are all positive

        Returns
        -------

        connection_probabilities: np.ndarray
            Connection probabilities
        """
        assert (positive_correlations >= 0).all(), "positive_correlations has to be all positive"
        # correlation to the binomial p
        p_binomial = positive_correlations ** self.n_pow  # this is a matrix
        connection_probabilities = sp.zeros((self.nr_neurons, self.nr_neurons))

        exc_population = sp.arange(self.nr_exc)
        inh_population = sp.arange(self.nr_exc, self.nr_neurons)
        for population, p_target in [(exc_population, self.p_exc), (inh_population, self.p_inh)]:
            binomial_fitter = BinomialConnectivity(p_binomial[:, population])
            binomial_fitter.fit(p_target)
            connection_probabilities[:, population] = binomial_fitter.predict()
            if population is exc_population:
                self.n_bin_exc_ = binomial_fitter.n_binomial_fitted_
            else:
                self.n_bin_inh_ = binomial_fitter.n_binomial_fitted_
        return connection_probabilities

    def signed_correlations_to_probabilities(self, signed_correlations):
        """

        Parameters
        ----------
        signed_correlations: np.ndarray
            Correlations between feature vectors. These are generally signed.

        Returns
        -------
        probabilities: np.ndarray
            Connection probabilities as calculated from correlations.
        """
        positive_correlations = self.signed_to_positive_correlations(signed_correlations)
        probabilities = self.positive_correlations_to_probability(positive_correlations)
        return probabilities

    def get_feature_vectors(self):
        return self.feature_vectors_

    def run(self):
        feature_vectors = uniform_spherical_feature_fectors(self.nr_neurons, self.feature_space_dimension)
        self.feature_vectors_ = feature_vectors

        correlations = feature_vector_correlations(feature_vectors)
        self.correlations_ = correlations

        signed_correlations = invert_sign_of_inh_to_exc_correlations(correlations, self.nr_inh)

        probabilities = self.signed_correlations_to_probabilities(signed_correlations)
        self.probabilities_ = probabilities

        network = sample_weighted_network_from_probabilities(probabilities, self.nr_exc, self.nr_inh)
        network.diag_to_zero()
        return network
