import sys
from .fever_lambda_optimizer import fever_lambda_optimizer
from .lasso_optimizer import DevelopmentalFeverLassoOptimizer
from modelflow import Input, Node
from .bases import NetworkModelMixin
from .er import ER
from .exp import EXP
from .featurevector import uniform_spherical_feature_fectors
import scipy as sp
import scipy.interpolate
import warnings
from abc import abstractmethod
FORCE_MONOTONIC_OPTIMIZER = True
from .network import Network


class NetworkTooSmallForLassoPath(Exception):
    pass


class FEVERGivenInitialGraph(Node):
    initial_network = Input()
    feature_vectors = Input()
    p_exc = Input()
    p_inh = Input()
    NR_NEURONS_TO_OPTIMIZE_FOR_LAMBDA_CALCULATION = 50

    def run(self):
        lambda_exc = self._lambda("exc")
        lambda_inh = self._lambda("inh")
        self.lambda_exc_ = lambda_exc
        self.lambda_inh_ = lambda_inh

        lasso = DevelopmentalFeverLassoOptimizer(nr_exc=self.initial_network.nr_exc)
        adjacency_matrix = lasso.fit(self.feature_vectors, self.initial_network.adjacency_matrix,
                                     lambda_exc, lambda_inh)

        network = self.initial_network.copy()
        network.adjacency_matrix = adjacency_matrix
        return network

    @staticmethod
    def _sanitize_adjacency(adjacency_matrix, nr_exc):
        if not (adjacency_matrix[:,:nr_exc] >= 0).all():
            nr_negative_exc_synapses = (~(adjacency_matrix[:,:nr_exc] >= 0)).sum()
            print("Warning: FEVER: Nr negative exc synapses: {}".format(nr_negative_exc_synapses), file=sys.stderr)
            adjacency_matrix[:,:nr_exc][~(adjacency_matrix[:,:nr_exc] >= 0)] = 0

        if not (adjacency_matrix[:,nr_exc:] <= 0).all():
            nr_positive_inh_synapses = (~(adjacency_matrix[:,nr_exc:] <= 0)).sum()
            print("Warning: FEVER: Nr positive inh synapses: {}".format(nr_positive_inh_synapses), file=sys.stderr)
            adjacency_matrix[:,nr_exc:][~(adjacency_matrix[:,nr_exc:] <= 0)] = 0

    def _lambda(self, population):
        """
        Calculate/optimize lambda to obtain desired population connectivity.

        Either use the LASSO path or use a monotonic optimizer to get the right lambda.

        Parameters
        ----------
        population: exc or inh
            For which population to optimize.

        Returns
        -------

        result: float
            Lambda matching the population connectivity.

        """
        if population == 'exc':
            p_target = self.p_exc
            neurons_to_optimize = range(0, min(self.NR_NEURONS_TO_OPTIMIZE_FOR_LAMBDA_CALCULATION,
                                               self.initial_network.nr_exc))
        elif population == 'inh':
            p_target = self.p_inh
            neurons_to_optimize = range(self.initial_network.nr_exc,
                                        min(self.initial_network.nr_exc
                                            + self.NR_NEURONS_TO_OPTIMIZE_FOR_LAMBDA_CALCULATION,
                                            self.initial_network.nr_neurons))
        lasso = DevelopmentalFeverLassoOptimizer(nr_exc=self.initial_network.nr_exc,
                                                 neurons_to_optimize=neurons_to_optimize)

        if FORCE_MONOTONIC_OPTIMIZER:
            return self._find_lambda_from_monotonic_optimization(p_target, lasso, neurons_to_optimize)
        try:
            return self._find_lambda_from_lasso_path(lasso, p_target)
        except NetworkTooSmallForLassoPath:
            warnings.warn("Could not determine lambda from LassoPath. Trying monotonic bisecting optimizer instead.")
            return self._find_lambda_from_monotonic_optimization(p_target, lasso, neurons_to_optimize)

    def _find_lambda_from_monotonic_optimization(self, p_target, lasso, neurons_to_optimize):
        def p_calculator(lmbda):
            adjacency_matrix = lasso.fit(self.feature_vectors, self.initial_network.adjacency_matrix, lmbda, lmbda)
            p = ((adjacency_matrix != 0)[:, neurons_to_optimize].sum()
                 / len(neurons_to_optimize)
                 / (self.initial_network.nr_neurons - 1))
            return p

        result = fever_lambda_optimizer(p_calculator, p_target)
        return result

    def _find_lambda_from_lasso_path(self, lasso, p_target):
        lasso_paths = lasso.fit(self.feature_vectors, self.initial_network.adjacency_matrix, -1, -1, lasso_path=True)
        connectivities_dict = {neuron: (alphas, (coefs != 0).sum(axis=0) / coefs.shape[0])
                          for neuron, (alphas, coefs) in lasso_paths.items()}
        if len(connectivities_dict) == 0:
            raise NetworkTooSmallForLassoPath()
        alpha_min = min(alpha.min() for alpha, coef in connectivities_dict.values())
        alpha_max = min(alpha.max() for alpha, coef in connectivities_dict.values())
        connectivity_of_alphas = [sp.interpolate.interp1d(alphas, connectivities, kind="linear", bounds_error=False,
                                             fill_value=(connectivities.min(), connectivities.max()))
                     for alphas, connectivities in connectivities_dict.values()]
        common_alphas = sp.logspace(sp.log10(alpha_min), sp.log10(alpha_max), 100)
        connectivities_at_common_alphas = sp.array([f(common_alphas) for f in connectivity_of_alphas])
        mean_connectivities = connectivities_at_common_alphas.mean(axis=0)

        valid_values = []
        current_lowest_connectivity = -1  # start with negativ connectivity to accept the last value
        for alpha, conn in zip(common_alphas[::-1], mean_connectivities[::-1]):
            if conn > current_lowest_connectivity:
                valid_values.append((alpha, conn))
                current_lowest_connectivity = conn
        valid_alphas, valid_connectivities = zip(*valid_values)
        lambda_of_connectivity = sp.interpolate.interp1d(valid_connectivities, valid_alphas,
                                                         kind="linear", bounds_error=False,
                                                         fill_value=(min(valid_alphas), max(valid_alphas)))
        result = lambda_of_connectivity(p_target)
        return result


class FEVERInitialConnectivity(NetworkModelMixin, Node):
    p_initial_graph_exc = Input()
    p_initial_graph_inh = Input()
    feature_space_dimension = Input()

    @abstractmethod
    def initial_network_model_factory(self) -> Node:
        pass

    def initial_network(self):
        initial_model = self.initial_network_model_factory()
        return initial_model(nr_neurons=self.nr_neurons, inh_ratio=self.inh_ratio,
                             p_exc=self.p_initial_graph_exc, p_inh=self.p_initial_graph_inh)

    def run(self) -> Network:
        feature_vectors = uniform_spherical_feature_fectors(self.nr_neurons, self.feature_space_dimension)
        self.feature_vectors_ = feature_vectors
        initial_network = self.initial_network()
        self.fever_given_initial_graph_ = FEVERGivenInitialGraph(initial_network=initial_network,
                                                          feature_vectors=feature_vectors,
                                                          p_exc=self.p_exc,
                                                          p_inh=self.p_inh)
        network = self.fever_given_initial_graph_()
        return network


class EXPFEVERInitialConnectivity(FEVERInitialConnectivity):
    def initial_network_model_factory(self):
        return EXP(decay=1)


class ERFEVERInitialConnectivity(FEVERInitialConnectivity):
    def initial_network_model_factory(self):
        return ER()


class FEVERFeverizationRatio(NetworkModelMixin, Node):
    """
    feverization_ratio is a value in the interval [0,1] determining how feverized
    the network is. a value of 0 correspond to the initial random graph.
    a value of 1 means it is maximally feverized, but still comes from the initial graph
    """
    FEVERInitialConnectivity = None  # A subclass of FEVERInitialConnectivity
    feverization_ratio = Input()
    feature_space_dimension = Input()

    def get_feature_vectors(self):
        return self.fever_initial_connectivity_.feature_vectors_

    def get_lambdas(self):
        lambda_exc = self.fever_initial_connectivity_.fever_given_initial_graph_.lambda_exc_
        lambda_inh = self.fever_initial_connectivity_.fever_given_initial_graph_.lambda_inh_
        return lambda_exc, lambda_inh

    def initial_connectivity(self, target_connectivity):
        """
        Calculate the necessary initial connectivity given the target connectivity
        """
        res = target_connectivity - self.feverization_ratio * self.feature_space_dimension / self.nr_neurons
        if res < 0:
            print("Warning: Initial connectivity was determined to be negative for FEVER. Set to 0 instead.",
                  file=sys.stderr)
        return max(0, res)

    def run(self):
        p_initial_graph_exc = self.initial_connectivity(self.p_exc)
        p_initial_graph_inh = self.initial_connectivity(self.p_inh)

        fever_initial_connectivity = self.FEVERInitialConnectivity()
        self.fever_initial_connectivity_ = fever_initial_connectivity
        network = fever_initial_connectivity(nr_neurons=self.nr_neurons, inh_ratio=self.inh_ratio,
                                  p_exc=self.p_exc, p_inh=self.p_inh,
                                  p_initial_graph_exc=p_initial_graph_exc,
                                  p_initial_graph_inh=p_initial_graph_inh,
                                  feature_space_dimension=self.feature_space_dimension)
        return network


class ERFEVER(FEVERFeverizationRatio):
    """
    FEVER network.

    Apart of the standard parameters of each :class:`network model <connectome.model.bases.NetworkModelMixin>`
    it has the following additional parameters:


    :param feature_space_dimension: int
        Dimension of the underlying feature space.
    :param feverization_ratio: float in the inverval [0,1]
        Determines to which degree the network interpolates
        between ER and FEVER. A value of feverization_ratio=0
        produces an ER graph. A value of 1, however, does *not*
        produce the original FEVER model. Instead, it produces
        a maximally feverized network, as determined by the
        connectivity and feature space dimension.
    """
    FEVERInitialConnectivity = ERFEVERInitialConnectivity


class EXPFEVER(FEVERFeverizationRatio):
    FEVERInitialConnectivity = EXPFEVERInitialConnectivity
