import sys
from collections import namedtuple

import numpy as np
import scipy as sp
from .block import Block

from .network import Network
from modelflow import Node
from modelflow.elementarynodes import Input
from .bases import connect_uniform, NetworkModelMixin

BlockConnectivityStructure = namedtuple("BlockConnectivityStructure", "blocks diagonal_connectivity off_diagonal_connectivity")


class NotUpperTriangular(Exception):
    pass


class UnsatisfiableConstraints(Exception):
    pass


def block_structure(nr_populations, nr_upper_diagonals):
    """
    Crates a binary matrix which represents the block structure

    Parameters
    ----------

    nr_upper_diagonals: int

    or 0 it is the diagonal only for nr_populations-1 it is complete uppper triangular
    """
    nr_populations = nr_populations
    ones = sp.ones((nr_populations, nr_populations))
    nr_upper_diagonals = nr_upper_diagonals
    blocks = sp.tril(sp.triu(ones), nr_upper_diagonals)
    return blocks


def add_inhibitory_adjacency(exc_adjacency, nr_inh):
    n_exc = exc_adjacency.shape[0]
    n_neurons = n_exc + nr_inh
    complete_adjacency = sp.zeros((n_neurons, n_neurons))
    complete_adjacency[:n_exc, :n_exc] = exc_adjacency
    return complete_adjacency


class BlockMatrix(Node):
    """
    Calculate connectivities by solving the formula

    .. math::
        p = p_l \\rho_l + p_f \\rho_ f

    and

    .. math::
        r = p(reciprocated connection|connection)

    .. math::
        p(reciprocated conn|conn) = p(rec conn | lateral conn) p(lateral conn| conn)

    And therfore

    .. math::
         r = p_l \\frac{p_l \\rho_l}{p}

    Where :math:`r` is the reciprocity,
    :math:`\\rho_l` is the number of diagonal blocks divided by the total number of blocks,
    :math:`\\rho_f` is the number of off diagonal blocks divided by the total number of blocks,
    and :math:`p_l` is the lateral (diagonal) connectivity and :math:`p_f` the forward (off diagonal)
    connectivity.
    """

    exc_connectivity = Input()
    exc_reciprocity = Input()
    nr_exc = Input()
    blocks = Input()  # boolean array like
    nr_inh = Input()

    def run(self) -> np.ndarray:
        block_structure = self.make_block_connectivities()
        exc_adjacency = self._block_structure_to_adjacency(block_structure)
        complete_adjacency = add_inhibitory_adjacency(exc_adjacency, self.nr_inh)
        blocks = []
        for k in range(self._nr_populations()):
            lower, upper = self._lower_upper_view_bound(k)
            size = upper-lower
            blocks.append(Block(size, "E"))
        blocks.append(Block(self.nr_inh, "I"))
        network = Network(complete_adjacency, nr_exc=exc_adjacency.shape[0], nr_inh=self.nr_inh, blocks=blocks)
        return network

    def make_block_connectivities(self):
        """
        Returns
        -------
        blocks, diagonal_connectivity, off_diagonal_connectivity
        """
        blocks = self._get_valid_block_structure()
        diagonal_connectivity, off_diagonal_connectivity = self._block_connectivities(blocks)
        self.diagonal_connectivity = diagonal_connectivity
        self.off_diagonal_connectivity_ = off_diagonal_connectivity
        return BlockConnectivityStructure(blocks, diagonal_connectivity, off_diagonal_connectivity)

    def _block_structure_to_adjacency(self, block_structure):
        blocks, diagonal_connectivity, off_diagonal_connectivity = block_structure
        adjacency_matrix = sp.zeros((self.nr_exc, self.nr_exc))
        for row in range(self._nr_populations()):
            for column in range(self._nr_populations()):
                if blocks[row, column] != 0:
                    self._fill_block(adjacency_matrix, row, column,
                                     diagonal_connectivity if row == column else off_diagonal_connectivity)
        return adjacency_matrix

    def _fill_block(self, adjacency_matrix, row, column, connectivity):
        view_on_block = self._get_view_on_block(adjacency_matrix, row, column)
        view_on_block[:,:] = (sp.rand(*view_on_block.shape) < connectivity).astype(float)

    def _get_view_on_block(self, adjacency_matrix, row, column):
        lower_row, upper_row = self._lower_upper_view_bound(row)
        lower_column, upper_column = self._lower_upper_view_bound(column)
        return adjacency_matrix[lower_row:upper_row, lower_column:upper_column]

    def _lower_upper_view_bound(self, population_nr):
        nr_populations = self._nr_populations()
        lower = int(sp.floor(population_nr / nr_populations * self.nr_exc))
        upper = int(sp.floor((population_nr + 1) / nr_populations * self.nr_exc))
        return lower, upper

    def _nr_populations(self):
        return sp.asarray(self.blocks).shape[0]

    def _block_connectivities(self, blocks):
        diagonal_block_density, off_diagonal_block_density = self._block_densities(blocks)
        diagonal_connectivity = self._diagonal_connectivity(diagonal_block_density)
        off_diagonal_connectivity = self._off_diagonal_connectivity(diagonal_block_density,
                                                                    diagonal_connectivity,
                                                                    off_diagonal_block_density)
        return diagonal_connectivity, off_diagonal_connectivity

    def _block_densities(self, blocks):
        max_nr_blocks = self._nr_populations()**2
        nr_diagonal_blocks = sp.diag(blocks).sum()
        nr_off_diagonal_blocks = sp.triu(blocks, 1).sum()
        diagonal_block_density = nr_diagonal_blocks / max_nr_blocks
        off_diagonal_block_density = nr_off_diagonal_blocks / max_nr_blocks
        return diagonal_block_density, off_diagonal_block_density

    def _off_diagonal_connectivity(self, diagonal_block_density, diagonal_connectivity,
                                   off_diagonal_block_density):
        if off_diagonal_block_density > 0:
            off_diagonal_connectivity = ((self.exc_connectivity - diagonal_connectivity * diagonal_block_density)
                                          / off_diagonal_block_density)
        else:
            if self.exc_reciprocity * diagonal_block_density != self.exc_connectivity:
                raise UnsatisfiableConstraints("Without off diagonal elements, constraints cannot be satisfied.")
            off_diagonal_connectivity = 0
        _check_connectivity_range(off_diagonal_connectivity, diagonal=False)
        return off_diagonal_connectivity

    def _diagonal_connectivity(self, diagonal_block_density):
        if diagonal_block_density > 0:
            diagonal_connectivity = sp.sqrt(self.exc_connectivity * self.exc_reciprocity / diagonal_block_density)
        else:
            diagonal_connectivity = 0
        _check_connectivity_range(diagonal_connectivity)
        return diagonal_connectivity

    def _get_valid_block_structure(self):
        blocks = sp.asarray(self.blocks)
        blocks = (blocks != 0).astype(int)
        if len(blocks.shape) != 2:
            raise ValueError("Require 2D block array")
        if blocks.shape[0] != blocks.shape[1]:
            raise ValueError("Require symmetric block array")
        if sp.tril(blocks, -1).sum() != 0:
            raise NotUpperTriangular("Invalid block structure")
        return blocks


def _check_connectivity_range(connectivity, diagonal=True):
    if not (0 <= connectivity <= 1):
        print("Block specification cannot be satisfied: {type} connectivity ({number}) ".format(
                number=connectivity,
                type="Diagonal" if diagonal else "Off diagonal")
            + "would not be in the interval [0, 1]", file=sys.stderr)


class BiologicalBlockmodel(NetworkModelMixin, Node):
    nr_exc_subpopulations = Input()
    reciprocity_exc = Input()
    nr_upper_diagonals = Input()

    def run(self):
        blocks = block_structure(self.nr_exc_subpopulations, self.nr_upper_diagonals)
        self.block_matrix_ = BlockMatrix(exc_connectivity=self.p_exc,
                                exc_reciprocity=self.reciprocity_exc,
                                nr_exc=self.nr_exc,
                                blocks=blocks,
                                nr_inh=self.nr_inh)
        network = self.block_matrix_()
        connect_uniform(network, "E", "I", self.p_exc)
        connect_uniform(network, "I", "E", self.p_inh)
        connect_uniform(network, "I", "I", self.p_inh)
        network.diag_to_zero()
        return network
