# -*- coding: utf-8 -*-
import scipy as sp
from sklearn import linear_model
from multiprocessing import Process, Queue
from parallel import nr_cores_available

ROW_NORMALIZATION = 1


def normalize_columns(initial_network: sp.ndarray, **pars):
    column_sum = sp.absolute(initial_network.sum(axis=0, keepdims=True))
    column_sum[column_sum == 0] = 1   # don't divide by zero
    initial_network /= column_sum


def normalize_rows(initial_network: sp.ndarray, **pars):
    nr_exc = pars["nr_exc"]
    nr_inh = initial_network.shape[1] - nr_exc
    exc_subpopulation = slice(0, initial_network.shape[0]), slice(0, nr_exc)
    inh_subpopulation = slice(0, initial_network.shape[0]), slice(nr_exc, initial_network.shape[1])

    normalize_subpopulation_row(initial_network, exc_subpopulation, normalize_to=ROW_NORMALIZATION)
    normalize_subpopulation_row(initial_network, inh_subpopulation, population_sign=-1,
                                normalize_to=ROW_NORMALIZATION*nr_exc/nr_inh if nr_inh != 0 else ROW_NORMALIZATION)


def normalize_subpopulation_row(initial_network: sp.ndarray, subpopulation, population_sign=1, normalize_to=1.):
    synapse_sum = initial_network[subpopulation].sum(axis=1, keepdims=True) * population_sign
    synapse_sum[synapse_sum == 0] = 1 # don't divide by zero
    initial_network[subpopulation] /= synapse_sum / normalize_to


class DevelopmentalFeverLassoOptimizer:
    """
    Optimize with respect to an initial bias.
    Use that the optimization problem
    ..math ::
        \\min_{W} \{ \\| D - DW\\|_2^2 + \\l \\|W_0 -W\\|_1 \\}.
    can be expressed as
    .. math ::
        \\min_{\\widetilde W} \\{ \\| D(id - W_0) - D \\widetilde W \\|_2^2 + \\l \\|\\widetilde W \\|_1 \\}.
    """
    def __init__(self, *, nr_exc, max_iter=100000, neurons_to_optimize=None,
                 initial_network_pre_processors=(normalize_columns, normalize_rows)):
        """
        Parameters
        ----------
        neurons_to_optimize:
            This is useful for getting connectivity statistics while not doing
            a full optimization of all neurons. Useful to calibrate lambda
            to achieve a desired connectivity.
            :param neurons_to_optimize: list of neurons to optimize
        """
        self.max_iter = max_iter
        self.nr_exc = nr_exc
        self.neurons_to_optimize = neurons_to_optimize
        self.initial_network_pre_processors = initial_network_pre_processors

    def fit(self, feature_vectors: sp.ndarray, initial_adjacency: sp.ndarray,
            lambda_exc: float, lambda_inh: float, lasso_path=False) -> sp.ndarray:
        """
        feature_vectors: ndarray
            feature matrix

        initial_adjacency : ndarray
            initial bias

        lambda_exc : float
            excitatory lamnda
            .. math::
             1/2 * \\| D - D result_feverized_weight_matrix \\| + lambda1 * n \\cdot \\| result_feverized_weight_matrix \\|_1

        lambda_inh: float
            inhibitory lambda

        Returns
        -------
          result_feverized_weight_matrix : ndarray
              the optimized weight matrix
        """
        if not (initial_adjacency.shape[0] == initial_adjacency.shape[1]):
            raise Exception("W0 has to be a square matrix.")
        if not (initial_adjacency.shape[0] == feature_vectors.shape[1]):
            raise Exception("Dimensions of W0 and D not matching")

        if self.neurons_to_optimize is not None:
            neurons_to_optimize = self.neurons_to_optimize
        else:
            neurons_to_optimize = list(range(initial_adjacency.shape[0]))


        # normalize
        for pre_processor in self.initial_network_pre_processors:
            pre_processor(initial_adjacency, nr_exc=self.nr_exc, nr_inh=initial_adjacency.shape[1] - self.nr_exc)

        # transform to ordinary lasso problem. See docstring of class.
        effective_target_vectors = self._make_effective_target_vectors(feature_vectors, initial_adjacency)
        result_feverized_weight_matrix = sp.zeros_like(initial_adjacency)


        # Use Queue and Processes instead of Pool.map
        # Uses Unix Fork and avoids copying large arrays
        # Therefore this Queue implementation is efficient
        task_queue = Queue()
        result_queue = Queue()
        nr_processes = nr_cores_available()
        processes = [OptimizationConsumer(task_queue, result_queue, self.nr_exc,
                                          effective_target_vectors, feature_vectors,
                                          initial_adjacency, lambda_exc, lambda_inh, self.max_iter,
                                          calculate_lasso_path=lasso_path)
                     for _ in range(nr_processes)]

        for process in processes:
            process.start()

        # If the number of neurons gets too large the queue might block
        # due to the Unix Pipe Buffer limit of 64K (64 kilobytes).
        # it might then be necessary to have another thread which feeds the queue
        for neuron in neurons_to_optimize:
            task_queue.put(neuron)

        # put -1 as sentinel to stop
        for _ in range(nr_processes):
            task_queue.put(-1)

        # collect results
        lasso_path_dict = {}
        for _ in range(len(neurons_to_optimize)):
            if not lasso_path:
                neuron, result_vector = result_queue.get()
                result_feverized_weight_matrix[:,neuron] = result_vector
            else:
                neuron, alphas, result_vectors = result_queue.get()
                lasso_path_dict[neuron] = (alphas, result_vectors)

        for process in processes:
            process.join()

        if not lasso_path:
            return result_feverized_weight_matrix
        else:
            return lasso_path_dict

    @staticmethod
    def _make_effective_target_vectors(feature_vectors, initial_network):
        """
        Transforms feature vectors to the X of ordinary the LASSO problem.

        Parameters
        ----------
        feature_vectors: ndarray of shape nr_features x nr_neurons
        initial_network: ndarray of shape nr_neurons x nr_neurons

        Returns
        -------

        effective_target_vectors: ndarray of shape nr_features x nr_neurons
        """
        effective_target_vectors = feature_vectors @ (sp.eye(initial_network.shape[0]) - initial_network)
        return effective_target_vectors


class OptimizationConsumer(Process):
    def __init__(self, task_queue: Queue, result_queue: Queue,
                 nr_exc, effective_target_vectors, feature_vectors,
                 initial_network, lambda_exc, lambda_inh, max_iter, calculate_lasso_path=False):
        super().__init__(daemon=True)
        self.task_queue = task_queue
        self.result_queue = result_queue
        self.feature_vectors = feature_vectors
        self.nr_exc = nr_exc
        self.effective_target_vectors = effective_target_vectors
        self.lambda_exc = lambda_exc
        self.lambda_inh = lambda_inh
        self.max_iter = max_iter
        self.initial_network = initial_network
        self.calculate_lasso_path = calculate_lasso_path

    def run(self):
        while True:
            neuron = self.task_queue.get()
            if neuron == -1:
                break

            sign, lmbda = (1, self.lambda_exc) if neuron < self.nr_exc else (-1, self.lambda_inh)
            own_target_vector = self.effective_target_vectors[:, neuron]
            # leave out vector neuron
            other_feature_vectors = self._make_other_feature_vectors(self.feature_vectors, neuron)
            # invert sign twice for negative neurons
            if not self.calculate_lasso_path:
                outgoing_weights_without_own = self._lasso(other_feature_vectors,
                                                               own_target_vector * sign, lmbda) * sign
                # reinsert zero at position k
                outgoing_weights_with_own = sp.insert(outgoing_weights_without_own, neuron, 0)
                result_vector = outgoing_weights_with_own + self.initial_network[:, neuron]

                self.result_queue.put((neuron, result_vector))
            else:
                alphas, outgoing_weights_without_own = self._lasso_path(other_feature_vectors,
                                                               own_target_vector * sign)
                outgoing_weights_without_own *= sign
                outgoing_weights_with_own = sp.insert(outgoing_weights_without_own, neuron, 0, axis=0)
                result_vectors = outgoing_weights_with_own + self.initial_network[:, neuron][:, sp.newaxis]
                self.result_queue.put((neuron, alphas, result_vectors))

    @staticmethod
    def _make_other_feature_vectors(feature_vectors, neuron_nr):
        other_feature_vectors = feature_vectors[:, list(range(neuron_nr))
                                                   + list(range(neuron_nr + 1, feature_vectors.shape[1]))]
        return other_feature_vectors

    def _lasso(self, dictionary, target, lambda1) -> sp.ndarray:
        try:
            assert sp.isfinite(lambda1), "Lambda has to be finite but got {}".format(lambda1)
        except TypeError:
            raise Exception("Colud not check lambda={} for finiteness".format(lambda1))

        assert sp.isfinite(dictionary).all(), "Dictionary has to be finite"
        assert sp.isfinite(target).all(), "Target has to be finite"
        skl_lasso = linear_model.Lasso(lambda1, positive=True, copy_X=True, max_iter=self.max_iter)
        # invert sign twice for negative neurons
        skl_lasso.fit(dictionary, target)
        return skl_lasso.coef_

    def _lasso_path(self, dictionary, target):
        alphas, coefs, _ = linear_model.lasso_path(dictionary, target, positive=True)
        return alphas, coefs
