import numpy as np 
import scipy as sp


def lattice(prng, n_features, alpha, random_sign=False, low=0.3, high=0.7):
    """Returns the adjacency matrix for a lattice network. 
    
    The resulting network is a Toeplitz matrix with random values summing 
    between -1 and 1 and zeros along the diagonal. 
    
    The range of the values can be controlled via the parameters low and high.  
    If random_sign is false, all entries will be negative, otherwise their sign
    will be modulated at random with probability 1/2.

    Each row has maximum edges of np.ceil(alpha * n_features).

    Parameters
    -----------        
    n_features : int 

    alpha : float (0, 1) 
        The complexity / sparsity factor.

    random sign : bool (default=False)
        Randomly modulate each entry by 1 or -1 with probability of 1/2.

    low : float (0, 1) (default=0.3)
        Lower bound for np.random.RandomState.uniform before normalization.

    high : float (0, 1) > low (default=0.7)
        Upper bound for np.random.RandomState.uniform before normalization.
    """
    adj = np.zeros((n_features, n_features))
    degree = int(1 + np.round(alpha * n_features / 2.))

    if random_sign:
        sign_row = (-1.0 * np.ones(degree) + 
                    2 * (prng.uniform(low=0, high=1, size=degree) > .5))
    else:
        sign_row = -1.0 * np.ones(degree)
    
    # in the *very unlikely* event that we draw a bad row that sums to zero
    # (which is only possible when random_sign=True), we try again up to 
    # MAX_ATTEMPTS=5 times.  If we are still unable to draw a good set of values
    # something is probably wrong and we raise.
    MAX_ATTEMPTS = 5
    attempt = 0
    row = np.zeros((n_features,))
    while np.sum(row) == 0 and attempt < MAX_ATTEMPTS:
        row = np.zeros((n_features,))
        row[1: 1 + degree] = sign_row * prng.uniform(low=low,
                                                     high=high,
                                                     size=degree)
        attempt += 1

    if np.sum(row) == 0:
        raise Exception('InvalidLattice', 'Rows sum to 0.')
        return

    # sum-normalize and keep signs
    row /= np.abs(np.sum(row))  

    return sp.linalg.toeplitz(c=row, r=row)


def blocks(prng, block, n_blocks=2, chain_blocks=True):
    """Replicates `block` matrix n_blocks times diagonally to create a
    square matrix of size n_features = block.size[0] * n_blocks and with zeros
    along the diagonal.

    The graph can be made fully connected using chaining assumption when 
    chain_blocks=True (default).
    
    This utility can be used to generate cluster networks or banded lattice n
    networks, among others.

    Parameters
    -----------        
    block : 2D array (n_block_features, n_block_features) 
        Prototype block adjacency matrix.

    n_blocks : int (default=2)
        Number of blocks.  Returned matrix will be square with 
        shape n_block_features * n_blocks.

    chain_blocks : bool (default=True)
        Apply random lattice structure to chain blocks together.
    """
    n_block_features, _ = block.shape
    n_features = n_block_features * n_blocks
    adjacency = np.zeros((n_features, n_features))    
        
    dep_groups = np.eye(n_blocks)
    if chain_blocks:
        chain_alpha = np.round(0.01 + 0.5 / n_blocks, 2)
        chain_groups = lattice(prng, n_blocks, chain_alpha, random_sign=False)
        chain_groups *= -0.1
        dep_groups += chain_groups
        
    adjacency = np.kron(dep_groups, block)
    adjacency[np.where(np.eye(n_features))] = 0 
    return adjacency


def _to_diagonally_dominant(mat):
    """Make matrix unweighted diagonally dominant using the Laplacian."""    
    mat += np.diag(np.sum(mat != 0, axis=1) + 0.01)
    return mat


def _to_diagonally_dominant_weighted(mat):
    """Make matrix weighted diagonally dominant using the Laplacian."""    
    mat += np.diag(np.sum(np.abs(mat), axis=1) + 0.01)
    return mat


def _rescale_to_unit_diagonals(mat):
    """Rescale matrix to have unit diagonals.

    Note: Call only after diagonal dominance is ensured.
    """   
    d = np.sqrt(np.diag(mat))
    mat /= d
    mat /= d[:, np.newaxis]
    return mat


class Graph(object):
    """Base class that returns the adjacency matrix for a network via
    the .create() method.

    Parameters
    ----------- 
    n_blocks : int (default=2)
        Number of blocks.  Returned matrix will be square with 
        shape n_block_features * n_blocks.

    chain_blocks : bool (default=True)
        Apply random lattice structure to chain blocks together.

    seed : int
        Seed for np.random.RandomState seed. (default=1)
    """
    def __init__(self, n_blocks=2, chain_blocks=True, seed=1):
        self.n_blocks = n_blocks
        self.chain_blocks = chain_blocks
        self.seed = seed
        self.prng = np.random.RandomState(self.seed)

    def to_precision(self, adjacency, weighted=True, rescale=True):
        if weighted:
            dd_adj = _to_diagonally_dominant_weighted(adjacency)
        else:
            dd_adj = _to_diagonally_dominant(adjacency)

        if rescale:
            return _rescale_to_unit_diagonals(dd_adj)

        return dd_adj

    def to_covariance(self, precision, rescale=True):
        covariance = np.linalg.inv(precision)
        
        if rescale:
            return _rescale_to_unit_diagonals(covariance)

        return covariance

    def prototype_adjacency(self, n_block_features, alpha):
        """Override this method with a custom base graph."""
        pass

    def create(self, n_features, alpha):
        """Build a new graph with block structure.

        Parameters
        -----------        
        n_features : int 

        alpha : float (0,1) 
            The complexity / sparsity factor for each graph type.
        
        Returns
        -----------  
        (n_features, n_features) matrices: covariance, precision, adjacency 
        """
        n_block_features = int(np.floor(1. * n_features / self.n_blocks))
        if n_block_features * self.n_blocks != n_features:
            raise ValueError(('Error: n_features {} not divisible by n_blocks {}.'
                              'Use n_features = n_blocks * int').format(
                            n_features,
                            self.n_blocks))
            return

        block_adj = self.prototype_adjacency(n_block_features, alpha)
        adjacency = blocks(self.prng,
                           block_adj,
                           n_blocks=self.n_blocks,
                           chain_blocks=self.chain_blocks)

        precision = self.to_precision(adjacency)
        covariance = self.to_covariance(precision)
        return covariance, precision, adjacency

@@ -4914,16 +4914,158 @@
 .seed)%0A%0A
+        if n_blocks == 1 and chain_blocks:%0A            raise ValueError(%22chain_blocks must be False when n_blocks is 1.%22)%0A            return%0A%0A
     def 

Add exception if graph is called with 1 block but chain_blocks is True.