import scipy as sp

from .util import diag_to_zero
from .block import Block, BlockStructure

class PopulationMismatch(Exception):
    pass


def save_normalize(arr, axis):
    """
    Normalize axis of arr s.t. arr.sum(axis) == 1 if the sum is != 0.
    If the sum is 0, everything is kept 0.
    Normalization is done in place.

    Parameters
    ----------
    arr: nump array
    axis: integer
    """
    axis_sum = arr.sum(axis, keepdims=True)
    axis_sum[axis_sum == 0] = 1
    arr /= axis_sum


class Network:
    """
    A thin wrapper around a numpy.ndarray.

    Supports excitatory and inhibitory neurons and performs additional sanity checks.

    Note: Diagonal entries are **allowed**.

    Network objects support views of the form network[start, end] where start, end are "E" or "I".

    If called with nr_exc and nr_inh a zeros filled adjacency matrix is created.



    Parameters
    ----------
    adjacency_matrix: array
        The adjacency matrix.

    nr_exc: optional int
        Number excitatory neurons.

    nr_inh: optional int
        Number excitatory neurons.

    blocks: optional :class:`BlockStructure <connectome.model.block.BlockStructure>`
        The block structure of the network.

    """
    def __init__(self, adjacency_matrix=None, nr_exc=None, nr_inh=None, blocks=None):

        if adjacency_matrix is not None:
            self._adjacency_matrix = sp.asarray(adjacency_matrix).copy()
        else:
            nr_exc = nr_exc or 0
            nr_inh = nr_inh or 0
            self._adjacency_matrix = sp.zeros((nr_exc+nr_inh, nr_exc+nr_inh))
        self.nr_exc = nr_exc or self._max_nr_exc(self._adjacency_matrix)
        if blocks:
            self.blocks = BlockStructure(blocks)
            assert sum(block.size for block in self.blocks) == self.nr_neurons, "Block sizes do not match network size"
        else:
            self.blocks = BlockStructure(self._infer_blocks())

    def __repr__(self):
        return "<Network nr_exc={nr_exc}, nr_inh={nr_inh}>".format(nr_exc=self.nr_exc, nr_inh=self.nr_inh)

    def _infer_blocks(self):
        return [Block(self.nr_exc, "E"), Block(self.nr_inh, "I")]

    @staticmethod
    def _max_nr_exc(matrix):
        try:
            return sp.amax(sp.argwhere(sp.all(matrix >= 0, axis=0))) + 1
        except ValueError:
            return 0

    @staticmethod
    def _max_nr_inh(matrix):
        try:
            amin = sp.amin(sp.argwhere(sp.all(matrix <= 0, axis=0)))
            return matrix.shape[1] - amin
        except ValueError:
            return 0

    @property
    def nr_neurons(self):
        "Total number of neurons in the network."
        return self._adjacency_matrix.shape[1]

    @property
    def nr_exc(self):
        "Number of excitatory neurons."
        return int(self._nr_exc)

    @nr_exc.setter
    def nr_exc(self, value):
        if value > self.nr_neurons:
            raise PopulationMismatch("Cannot set number excitatory neurons to " + str(value) + " since this is larger than"
                            "number neurons (" + str(self.nr_neurons) + ")")
        if value > self._max_nr_exc(self._adjacency_matrix):
            raise PopulationMismatch("Nr excitatory too large. There are neurons with negative synapses.")
        if value < 0:
            raise PopulationMismatch("Nr excitatory neurons has to be greater than zero.")
        self._nr_exc = value

    @property
    def nr_inh(self):
        "Number of inhibitory neurons."
        return int(self.nr_neurons - self.nr_exc)

    @nr_inh.setter
    def nr_inh(self, nr_inh):
        self.nr_exc = self.nr_neurons - nr_inh

    @property
    def adjacency_matrix(self):
        "A view on the underlying adjacency matrix."
        return self._adjacency_matrix

    @property
    def shape(self):
        "The shape of the underlying adjacency matrix."
        return self._adjacency_matrix.shape

    @adjacency_matrix.setter
    def adjacency_matrix(self, matrix):
        if self.shape != matrix.shape:
            raise PopulationMismatch("Invalid shape: " + str(matrix.shape) + ". Should be " + str(self.shape))
        if (self._max_nr_exc(matrix) < self.nr_exc) or (self._max_nr_inh(matrix) < self.nr_inh):
            raise PopulationMismatch("Matrix is incompatible with number excitatory and inhibitory neurons: "
                                     + "n_exc={}, n_inh={} but new matrix allows maximally n_exc={}, n_inh={}"
                                     .format(self.nr_exc, self.nr_inh, self._max_nr_exc(matrix), self._max_nr_inh(matrix)))
        self._adjacency_matrix = matrix.copy()

    def _view(self, pre=None, post=None):
        if pre is post is None:
            return self._adjacency_matrix[:,:]
        if pre == 'E' and post == 'E':
            return self._adjacency_matrix[:self.nr_exc, :self.nr_exc]
        if pre == 'E' and post == 'I':
            return self._adjacency_matrix[self.nr_exc:, :self.nr_exc]
        if pre == 'I' and post == 'I':
            return self._adjacency_matrix[self.nr_exc:, self.nr_exc:]
        if pre == 'I' and post == 'E':
            return self._adjacency_matrix[:self.nr_exc, self.nr_exc:]

        if pre == "E":
            return self._adjacency_matrix[post, :self.nr_exc]
        if post == "E":
            return self._adjacency_matrix[:self.nr_exc, pre]
        if pre == "I":
            return self._adjacency_matrix[post, self.nr_exc:]
        if post == "I":
            return self._adjacency_matrix[self.nr_exc:, pre]

    def __getitem__(self, item):
        try:
            return self._view(item[0].start, item[1].start)
        except AttributeError:
            return self._view(*item)

    def __setitem__(self, key, value):
        view = self[key]
        view[:,:] = value

    @property
    def shape(self):
        return self._adjacency_matrix.shape

    def __imatmul__(self, other):
        self._adjacency_matrix @= other
        return self

    def __imul__(self, other):
        self._adjacency_matrix *= other
        return self

    def __itruediv__(self, other):
        self._adjacency_matrix /= other
        return self

    def normalize(self, in_out="in", exc=1, inh=1):
        """
        Normalize incoming or outgoing weights.

        Parameters
        ----------
        n_out: str, "in" or "out"
            Whether to normalize the incoming or outgoing weights.
        exc: float
            The weight to which to normalize excitatory neurons.
        inh: float
            The weight to which to normalize inhibitory neurons.
        """
        self._adjacency_matrix = self._adjacency_matrix.astype(float)
        if in_out == "in":
            axis = 1
        elif in_out == "out":
            axis = 0
        else:
            raise Exception("in_out can either be \"in\" our \"out\"")

        for sign, source, value in [(1, "E", exc), (-1, "I", inh)]:
            subnet = self[source, :]
            save_normalize(subnet, axis)
            subnet *= sign * value

    def _generate_new_labels(self):
        """

        Returns
        -------
        new labels: tuple
            The tuple is of the form (new_exc_label, new_inh_labels).
        """
        return sp.random.choice(self.nr_exc, self.nr_exc, False),  sp.random.choice(self.nr_inh, self.nr_inh, False)

    def random_relabel_neurons(self):
        """
        Randomly relabel all neurons.
        """
        new_exc_labels, new_inh_labels = self._generate_new_labels()
        self._relabel_exc_neurons(new_exc_labels)
        self._relabel_inh_neurons(new_inh_labels)
        self.blocks = BlockStructure(self._infer_blocks())

    def _relabel_exc_neurons(self, new_exc_labels):
        W = self._adjacency_matrix
        W[:, :self.nr_exc] = W[:, :self.nr_exc][:, new_exc_labels]
        W[:self.nr_exc, :] = W[:self.nr_exc, :][new_exc_labels, :]

    def _relabel_inh_neurons(self, new_inh_labels):
        W = self._adjacency_matrix
        W[:, self.nr_exc:] = W[:, self.nr_exc:][:, new_inh_labels]
        W[self.nr_exc:, :] = W[self.nr_exc:, :][new_inh_labels, :]

    def binarize(self):
        """
        Set any nonzero excitatory weight to 1 and any nonzero inhibitory weight to -1.
        """
        self._adjacency_matrix[self._adjacency_matrix > 0] = 1
        self._adjacency_matrix[self._adjacency_matrix < 0] = -1

    def diag_to_zero(self):
        diag_to_zero(self._adjacency_matrix)

    def copy(self):
        """
        Returns
        -------
        Network
            A copy of the network.
        """
        return self.__class__(self.adjacency_matrix.copy(), nr_exc=self.nr_exc, nr_inh=self.nr_inh)

    def __eq__(self, other):
        if self.nr_exc != other.nr_exc:
            return False
        if self.nr_inh != other.nr_inh:
            return False
        if (self.adjacency_matrix != other.adjacency_matrix).any():
            return False
        return True

    def _neurons_to_keep(self, neuron_numbers):
        return sp.asarray([k for k in range(self.nr_neurons) if k not in neuron_numbers]).astype(int)

    def remove_neurons(self, neuron_numbers):
        """
        Remove neurons from the network.

        Parameters
        ----------
        neuron_numbers: list
            A list of the neuron indices which to remove.
        """
        neuron_numbers = sp.asarray(neuron_numbers)
        nr_exc_to_remove = (neuron_numbers < self.nr_exc).sum()
        neurons_to_keep = self._neurons_to_keep(neuron_numbers)[sp.newaxis]
        self._adjacency_matrix = self._adjacency_matrix[neurons_to_keep.T, neurons_to_keep]
        self._nr_exc -= nr_exc_to_remove

    def blocks(self):
        raise NotImplementedError()

    def get_block(self, *, target_ix, source_ix):
        """
        Return a view on a block.

        Parameters
        ----------

        target_ix: int
            Target block number.
        source_ix: int
            Source block number.

        Returns
        -------
        view_on_subarray: ndarray
            The view on the subarray.
        """
        target_start, target_end = self.blocks.start_end(target_ix)
        source_start, source_end = self.blocks.start_end(source_ix)
        return self._adjacency_matrix[target_start:target_end, source_start:source_end]



class EmbeddedNetwork(Network):
    def __init__(self, adjacency_matrix: "array like"=None, nr_exc=None, nr_inh=None):
        super().__init__(adjacency_matrix=adjacency_matrix, nr_exc=nr_exc, nr_inh=nr_inh)
        self._positions = sp.zeros((self.nr_neurons, 3))

    @property
    def positions(self):
        return self._positions

    @positions.setter
    def positions(self, value):
        if self.shape[0] != value.shape[0]:
            raise PopulationMismatch("Need as many positions as neurons.")
        self._positions = value

    def _relabel_exc_neurons(self, new_exc_labels):
        super()._relabel_exc_neurons(new_exc_labels)
        self._positions[:self.nr_exc] = self._positions[:self.nr_exc][new_exc_labels]

    def _relabel_inh_neurons(self, new_inh_labels):
        super()._relabel_inh_neurons(new_inh_labels)
        self._positions[self.nr_exc:] = self._positions[self.nr_exc:][new_inh_labels]

    def copy(self):
        network = super().copy()
        network.positions = self.positions.copy()
        return network

    def __eq__(self, other):
        super().__eq__(other)
        try:
            if (self.positions != other.positions).any():
                return False
        except AttributeError:
            return False
        return True

    def remove_neurons(self, neuron_numbers):
        neurons_to_keep = self._neurons_to_keep(neuron_numbers)
        super().remove_neurons(neuron_numbers)
        positions = self.positions[neurons_to_keep]
        self._positions = positions
