# Path: pyphi/distribution.py
# def normalize(a):
# def uniform_distribution(number_of_nodes):
# def marginal_zero(repertoire, node_index):
# def marginal(repertoire, node_index):
# def independent(repertoire):
# def purview(repertoire):
# def purview_size(repertoire):
# def repertoire_shape(purview, N):  # pylint: disable=redefined-outer-name
# def flatten(repertoire, big_endian=False):
# def unflatten(repertoire, purview, N, big_endian=False):
# def max_entropy_distribution(node_indices, number_of_nodes):
# 
# Path: pyphi/utils.py
# def powerset(iterable, nonempty=False, reverse=False):
#     """Generate the power set of an iterable.
# 
#     Args:
#         iterable (Iterable): The iterable from which to generate the power set.
# 
#     Keyword Args:
#         nonempty (boolean): If True, don't include the empty set.
#         reverse (boolean): If True, reverse the order of the powerset.
# 
#     Returns:
#         Iterable: An iterator over the power set.
# 
#     Example:
#         >>> ps = powerset(np.arange(2))
#         >>> list(ps)
#         [(), (0,), (1,), (0, 1)]
#         >>> ps = powerset(np.arange(2), nonempty=True)
#         >>> list(ps)
#         [(0,), (1,), (0, 1)]
#         >>> ps = powerset(np.arange(2), nonempty=True, reverse=True)
#         >>> list(ps)
#         [(1, 0), (1,), (0,)]
#     """
#     iterable = list(iterable)
# 
#     if nonempty:  # Don't include 0-length subsets
#         start = 1
#     else:
#         start = 0
# 
#     seq_sizes = range(start, len(iterable) + 1)
# 
#     if reverse:
#         seq_sizes = reversed(seq_sizes)
#         iterable.reverse()
# 
#     return chain.from_iterable(combinations(iterable, r) for r in seq_sizes)


import numpy as np
from pyphi import distribution
from pyphi.utils import powerset


def test_flatten():
    # fmt: off
    repertoire = np.array([
        [[0.1, 0.0]],
        [[0.2, 0.7]],
    ])
    # fmt: on

    assert np.array_equal(distribution.flatten(repertoire), [0.1, 0.2, 0.0, 0.7])
    assert np.array_equal(
        distribution.flatten(repertoire, big_endian=True), [0.1, 0.0, 0.2, 0.7]
    )
    assert distribution.flatten(None) is None


def test_unflatten():
    repertoire = np.array(
        [
            [[0.1, 0.0]],
            [[0.2, 0.7]],
        ]
    )
    flat_repertoire = np.array([0.1, 0.2, 0.0, 0.7])
    assert distribution.unflatten(flat_repertoire, (0, 1), 2).shape == (2, 2)
    assert np.array_equal(
        distribution.unflatten(flat_repertoire, (0, 2), 3),
        repertoire,
    )

# Path: pyphi/distribution.py
# def normalize(a):
# def uniform_distribution(number_of_nodes):
# def marginal_zero(repertoire, node_index):
# def marginal(repertoire, node_index):
# def independent(repertoire):
# def purview(repertoire):
# def purview_size(repertoire):
# def repertoire_shape(purview, N):  # pylint: disable=redefined-outer-name
# def flatten(repertoire, big_endian=False):
# def unflatten(repertoire, purview, N, big_endian=False):
# def max_entropy_distribution(node_indices, number_of_nodes):
# 
# Path: pyphi/utils.py
# def powerset(iterable, nonempty=False, reverse=False):
#     """Generate the power set of an iterable.
# 
#     Args:
#         iterable (Iterable): The iterable from which to generate the power set.
# 
#     Keyword Args:
#         nonempty (boolean): If True, don't include the empty set.
#         reverse (boolean): If True, reverse the order of the powerset.
# 
#     Returns:
#         Iterable: An iterator over the power set.
# 
#     Example:
#         >>> ps = powerset(np.arange(2))
#         >>> list(ps)
#         [(), (0,), (1,), (0, 1)]
#         >>> ps = powerset(np.arange(2), nonempty=True)
#         >>> list(ps)
#         [(0,), (1,), (0, 1)]
#         >>> ps = powerset(np.arange(2), nonempty=True, reverse=True)
#         >>> list(ps)
#         [(1, 0), (1,), (0,)]
#     """
#     iterable = list(iterable)
# 
#     if nonempty:  # Don't include 0-length subsets
#         start = 1
#     else:
#         start = 0
# 
#     seq_sizes = range(start, len(iterable) + 1)
# 
#     if reverse:
#         seq_sizes = reversed(seq_sizes)
#         iterable.reverse()
# 
#     return chain.from_iterable(combinations(iterable, r) for r in seq_sizes)


# Path: test/test_distribution.py
import numpy as np
from pyphi import distribution
from pyphi.utils import powerset



def test_flatten():
    # fmt: off
    repertoire = np.array([
        [[0.1, 0.0]],
        [[0.2, 0.7]],
    ])
    # fmt: on

    assert np.array_equal(distribution.flatten(repertoire), [0.1, 0.2, 0.0, 0.7])
    assert np.array_equal(
        distribution.flatten(repertoire, big_endian=True), [0.1, 0.0, 0.2, 0.7]
    )
    assert distribution.flatten(None) is None


def test_unflatten():
    repertoire = np.array(
        [
            [[0.1, 0.0]],
            [[0.2, 0.7]],
        ]
    )
    flat_repertoire = np.array([0.1, 0.2, 0.0, 0.7])
    assert distribution.unflatten(flat_repertoire, (0, 1), 2).shape == (2, 2)
    assert np.array_equal(
        distribution.unflatten(flat_repertoire, (0, 2), 3),
        repertoire,
    )

    assert distribution.unflatten(flat_repertoire, (2, 5), 8).shape == (